Unverified Commit ceca097f authored by Asit's avatar Asit Committed by GitHub
Browse files

editing comments

parent 02a33875
...@@ -6,7 +6,7 @@ torchvision_imported=True ...@@ -6,7 +6,7 @@ torchvision_imported=True
try: try:
import torchvision import torchvision
except ImportError: except ImportError:
print("[ASP][Warning] torchvision cannot be imported, may infuence functionality of MaskRCNN/KeypointRCNN network from torchvision.") print("[ASP][Warning] torchvision cannot be imported.")
torchvision_imported=False torchvision_imported=False
def eligible_modules(model, whitelist_layer_types, allowed_layer_names, disallowed_layer_names): def eligible_modules(model, whitelist_layer_types, allowed_layer_names, disallowed_layer_names):
...@@ -78,7 +78,7 @@ class ASP: ...@@ -78,7 +78,7 @@ class ASP:
# function to extract variables that will be sparsified. # function to extract variables that will be sparsified.
# idea is that you will add one of these functions for each module type that can be sparsified. # idea is that you will add one of these functions for each module type that can be sparsified.
if torchvision_imported: if torchvision_imported:
print("[ASP] torchvision is imported, can work smoothly with the MaskRCNN/KeypointRCNN from torchvision.") print("[ASP] torchvision is imported, can work with the MaskRCNN/KeypointRCNN from torchvision.")
sparse_parameter_list = {torch.nn.Linear: ['weight'], torch.nn.Conv1d: ['weight'], torch.nn.Conv2d: ['weight'], torch.nn.Conv3d: ['weight'], torchvision.ops.misc.Conv2d: ['weight']} sparse_parameter_list = {torch.nn.Linear: ['weight'], torch.nn.Conv1d: ['weight'], torch.nn.Conv2d: ['weight'], torch.nn.Conv3d: ['weight'], torchvision.ops.misc.Conv2d: ['weight']}
else: else:
sparse_parameter_list = {torch.nn.Linear: ['weight'], torch.nn.Conv1d: ['weight'], torch.nn.Conv2d: ['weight'], torch.nn.Conv3d: ['weight']} sparse_parameter_list = {torch.nn.Linear: ['weight'], torch.nn.Conv1d: ['weight'], torch.nn.Conv2d: ['weight'], torch.nn.Conv3d: ['weight']}
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment