[models] Basemodel returns a dict intead of a list.

Author: arthurdouillard

inclearn/lib/network/basenet.py

@@ -100,19 +100,18 @@ def forward(self, x):
         selected_outputs = outputs[0] if self.classifier_no_act else outputs[1]
         logits = self.classifier(self.dropout(selected_outputs))
 
+        outputs = {"logits": logits}
+
         if self.return_features:
-            to_return = []
             if self.extract_no_act:
-                to_return.append(outputs[0])
+                outputs["features"] = outputs[0]
             else:
-                to_return.append(outputs[1])
+                outputs["features"] =  outputs[1]
 
-            to_return.append(logits)
             if self.attention_hook:
-                to_return.append(outputs[2])
+                outputs["attention_maps"] = outputs[2]
 
-            return to_return
-        return logits
+        return outputs
 
     def post_process(self, x):
         if self.post_processor is None:

inclearn/models/bic.py

@@ -82,11 +82,14 @@ def _eval_task(self, loader):
 
         return ypred, ytrue
 
-    def _compute_loss(self, inputs, logits, targets, onehot_targets, memory_flags):
+    def _compute_loss(self, inputs, outputs, targets, onehot_targets, memory_flags):
+        logits = outputs["logits"]
+
         loss = F.cross_entropy(logits, targets)
 
         if self._old_model is not None:
-            old_targets = self._old_model.post_process(self._old_model(inputs)).detach()
+            with torch.no_grad():
+                old_targets = self._old_model.post_process(self._old_model(inputs)["logits"])
 
             loss += F.binary_cross_entropy_with_logits(
                 logits[..., :-self._task_size] / self._temperature,

inclearn/models/icarl.py

@@ -233,9 +233,9 @@ def _forward_loss(self, training_network, inputs, targets, memory_flags):
         inputs, targets = inputs.to(self._device), targets.to(self._device)
         onehot_targets = utils.to_onehot(targets, self._n_classes).to(self._device)
 
-        logits = training_network(inputs)
+        outputs = training_network(inputs)
 
-        loss = self._compute_loss(inputs, logits, targets, onehot_targets, memory_flags)
+        loss = self._compute_loss(inputs, outputs, targets, onehot_targets, memory_flags)
 
         if not utils._check_loss(loss):
             pdb.set_trace()
@@ -270,11 +270,14 @@ def _eval_task(self, data_loader):
     # Private API
     # -----------
 
-    def _compute_loss(self, inputs, logits, targets, onehot_targets, memory_flags):
+    def _compute_loss(self, inputs, outputs, targets, onehot_targets, memory_flags):
+        logits = outputs["logits"]
+
         if self._old_model is None:
             loss = F.binary_cross_entropy_with_logits(logits, onehot_targets)
         else:
-            old_targets = torch.sigmoid(self._old_model(inputs).detach())
+            with torch.no_grad():
+                old_targets = torch.sigmoid(self._old_model(inputs)["outputs"])
 
             new_targets = onehot_targets.clone()
             new_targets[..., :-self._task_size] = old_targets

inclearn/models/still.py

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