update unit test

fangyixiao18 · fangyixiao18 · commit a98a31442e44 · 2022-03-11T14:16:59.000+08:00
diff --git a/mmselfsup/models/algorithms/base.py b/mmselfsup/models/algorithms/base.py
@@ -150,7 +150,10 @@ def val_step(self, data, optimizer):
         losses = self(**data)
         loss, log_vars = self._parse_losses(losses)
 
-        outputs = dict(
-            loss=loss, log_vars=log_vars, num_samples=len(data['img'].data))
+        if isinstance(data['img'], list):
+            num_samples = len(data['img'][0].data)
+        else:
+            num_samples = len(data['img'].data)
+        outputs = dict(loss=loss, log_vars=log_vars, num_samples=num_samples)
 
         return outputs
diff --git a/tests/test_models/test_algorithms/test_deepcluster.py b/tests/test_models/test_algorithms/test_deepcluster.py
@@ -36,7 +36,7 @@ def test_deepcluster():
 
     fake_input = torch.randn((16, 3, 224, 224))
     fake_labels = torch.ones(16, dtype=torch.long)
-    fake_out = alg.forward_test(fake_input)
+    fake_out = alg.forward(fake_input, mode='test')
     assert 'head0' in fake_out
     assert fake_out['head0'].size() == torch.Size([16, num_classes])
 
diff --git a/tests/test_models/test_algorithms/test_densecl.py b/tests/test_models/test_algorithms/test_densecl.py
@@ -77,3 +77,8 @@ def test_densecl():
     assert fake_loss['loss_dense'] > 0
     assert alg.queue_ptr.item() == 16
     assert alg.queue2_ptr.item() == 16
+
+    # test train step with 2 keys in loss
+    fake_outputs = alg.train_step(dict(img=[fake_input, fake_input]), None)
+    assert fake_outputs['loss'].item() > -1
+    assert fake_outputs['num_samples'] == 16
diff --git a/tests/test_models/test_algorithms/test_mocov3.py b/tests/test_models/test_algorithms/test_mocov3.py
@@ -52,6 +52,6 @@ def test_mocov3():
     alg.momentum_update()
 
     fake_input = torch.randn((16, 3, 224, 224))
-    fake_backbone_out = alg.extract_feat(fake_input)
+    fake_backbone_out = alg.forward(fake_input, mode='extract')
     assert fake_backbone_out[0][0].size() == torch.Size([16, 384, 14, 14])
     assert fake_backbone_out[0][1].size() == torch.Size([16, 384])
diff --git a/tests/test_models/test_algorithms/test_simsiam.py b/tests/test_models/test_algorithms/test_simsiam.py
@@ -51,3 +51,13 @@ def test_simsiam():
     ]
     fake_out = alg.forward(fake_input)
     assert fake_out['loss'].item() > -1
+
+    # test train step
+    fake_outputs = alg.train_step(dict(img=fake_input), None)
+    assert fake_outputs['loss'].item() > -1
+    assert fake_outputs['num_samples'] == 16
+
+    # test val step
+    fake_outputs = alg.val_step(dict(img=fake_input), None)
+    assert fake_outputs['loss'].item() > -1
+    assert fake_outputs['num_samples'] == 16
