[convnet] Returns a dictionnary instead of list;

Author: arthurdouillard

inclearn/convnet/my_resnet.py

@@ -261,7 +261,11 @@ def _make_layer(self, Block, planes, increase_dim=False, n=None):
 
         return Stage(layers, block_relu=self.last_relu)
 
-    def forward(self, x, attention_hook=False):
+    @property
+    def last_conv(self):
+        return self.stage_4.conv_b
+
+    def forward(self, x):
         x = self.conv_1_3x3(x)
         x = F.relu(self.bn_1(x), inplace=True)
 
@@ -278,9 +282,11 @@ def forward(self, x, attention_hook=False):
         else:
             attentions = [feats_s1[-1], feats_s2[-1], feats_s3[-1], x]
 
-        if attention_hook:
-            return raw_features, features, attentions
-        return raw_features, features
+        return {
+            "raw_features": raw_features,
+            "features": features,
+            "attention": attentions
+        }
 
     def end_features(self, x):
         x = self.pool(x)

inclearn/convnet/resnet.py

@@ -159,7 +159,11 @@ def _make_layer(self, block, planes, blocks, stride=1, last=False):
 
         return nn.Sequential(*layers)
 
-    def forward(self, x, attention_hook=False):
+    @property
+    def last_conv(self):
+        return self.layer4[-1].conv2
+
+    def forward(self, x):
         x = self.conv1(x)
         x = self.bn1(x)
         x = self.relu(x)
@@ -173,9 +177,11 @@ def forward(self, x, attention_hook=False):
         raw_features = self.end_features(x_4)
         features = self.end_features(F.relu(x_4, inplace=False))
 
-        if attention_hook:
-            return raw_features, features, [x_1, x_2, x_3, x_4]
-        return raw_features, features
+        return {
+            "raw_features": raw_features,
+            "features": features,
+            "attention": [x_1, x_2, x_3, x_4]
+        }
 
     def end_features(self, x):
         x = self.avgpool(x)

inclearn/models/still.py

@@ -352,7 +352,7 @@ def _before_task(self, train_loader, val_loader):
             self._class_weights = None
 
     def _compute_loss(self, inputs, outputs, targets, onehot_targets, memory_flags):
-        features, logits, atts = outputs["features"], outputs["logits"], outputs["attention_maps"]
+        features, logits, atts = outputs["raw_features"], outputs["logits"], outputs["attention"]
 
         if self._post_processing_type is None:
             scaled_logits = self._network.post_process(logits)
@@ -362,8 +362,8 @@ def _compute_loss(self, inputs, outputs, targets, onehot_targets, memory_flags):
         if self._old_model is not None:
             with torch.no_grad():
                 old_outputs = self._old_model(inputs)
-                old_features = old_outputs["features"]
-                old_atts = old_outputs["attention_maps"]
+                old_features = old_outputs["raw_features"]
+                old_atts = old_outputs["attention"]
 
         if self._ams_config:
             ams_config = copy.deepcopy(self._ams_config)