don't extract BN layers for list of weight matrices & wrote test

Author: McCrearyD

rigl_torch/util.py

@@ -2,6 +2,9 @@
 import torchvision
 
 
+EXCLUDED_TYPES = (torch.nn.BatchNorm2d, )
+
+
 def get_weighted_layers(model, i=0, layers=None, linear_layers_mask=None):
     if layers is None:
         layers = []
@@ -16,18 +19,9 @@ def get_weighted_layers(model, i=0, layers=None, linear_layers_mask=None):
         if isinstance(p, torch.nn.Linear):
             layers.append([p])
             linear_layers_mask.append(1)
-        elif hasattr(p, 'weight'):
+        elif hasattr(p, 'weight') and type(p) not in EXCLUDED_TYPES:
             layers.append([p])
             linear_layers_mask.append(0)
-        elif isinstance(p, torch.nn.BatchNorm2d):
-            layers[-1].append(p)
-        elif isinstance(p, torch.nn.ReLU):
-            layers[-1].append(p)
-        elif isinstance(p, torch.nn.MaxPool2d) or isinstance(p, torch.nn.AdaptiveAvgPool2d):
-            layers[-1].append(p)
-        elif layer_name == 'downsample':
-            layers.append(p)
-            linear_layers_mask.append(0)
         elif isinstance(p, torchvision.models.resnet.Bottleneck) or isinstance(p, torchvision.models.resnet.BasicBlock):
             _, linear_layers_mask, i = get_weighted_layers(p, i=i + 1, layers=layers, linear_layers_mask=linear_layers_mask)
         else:

tests/test_rigl.py

@@ -24,6 +24,16 @@
 criterion = torch.nn.functional.cross_entropy
 
 
+def test_lengths_of_W():
+    resnet18 = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=False)
+    resnet18_W = get_W(resnet18)
+    assert len(resnet18_W) == 21, 'resnet18 should have 21 "weight" matrices'
+
+    resnet50 = torch.hub.load('pytorch/vision:v0.6.0', 'resnet50', pretrained=False)
+    resnet50_W = get_W(resnet50)
+    assert len(resnet50_W) == 54, 'resnet50 should have 54 "weight" matrices'
+
+
 def get_dummy_dataloader():
     X = torch.rand((max_iters, *image_dimensionality))
     T = (torch.rand(max_iters) * num_classes).long()