From 955632e567e9166152a9c31bc4223ea31670ccc1 Mon Sep 17 00:00:00 2001 From: Yixiao Fang <36138628+fangyixiao18@users.noreply.github.com> Date: Thu, 28 Apr 2022 00:43:08 +0800 Subject: [PATCH] [UT] reduce memory usage while runing unit test (#291) --- .../test_models/test_algorithms/test_byol.py | 23 +++++++------- .../test_algorithms/test_classification.py | 14 ++++----- .../test_algorithms/test_deepcluster.py | 10 +++---- .../test_algorithms/test_densecl.py | 20 ++++++------- tests/test_models/test_algorithms/test_mae.py | 4 +-- .../test_mmcls_classifier_wrapper.py | 8 ++--- .../test_models/test_algorithms/test_moco.py | 16 +++++----- .../test_algorithms/test_mocov3.py | 16 +++++----- .../test_models/test_algorithms/test_npid.py | 10 +++---- tests/test_models/test_algorithms/test_odc.py | 16 +++++----- .../test_algorithms/test_relative_loc.py | 12 ++++---- .../test_algorithms/test_rotation_pred.py | 9 +++--- .../test_algorithms/test_simclr.py | 12 ++++---- .../test_algorithms/test_simsiam.py | 25 +++++++--------- .../test_models/test_algorithms/test_swav.py | 30 +++++++++---------- 15 files changed, 109 insertions(+), 116 deletions(-) diff --git a/tests/test_models/test_algorithms/test_byol.py b/tests/test_models/test_algorithms/test_byol.py index 7a34a3e98..c0d74d5bb 100644 --- a/tests/test_models/test_algorithms/test_byol.py +++ b/tests/test_models/test_algorithms/test_byol.py @@ -8,15 +8,15 @@ backbone = dict( type='ResNet', - depth=50, + depth=18, in_channels=3, out_indices=[4], # 0: conv-1, x: stage-x norm_cfg=dict(type='BN')) neck = dict( type='NonLinearNeck', - in_channels=2048, - hid_channels=4, - out_channels=4, + in_channels=512, + hid_channels=2, + out_channels=2, with_bias=True, with_last_bn=False, with_avg_pool=True, @@ -25,9 +25,9 @@ type='LatentPredictHead', predictor=dict( type='NonLinearNeck', - in_channels=4, - hid_channels=4, - out_channels=4, + in_channels=2, + hid_channels=2, + out_channels=2, with_bias=True, with_last_bn=False, with_avg_pool=False, @@ -42,15 +42,12 @@ def test_byol(): alg = BYOL(backbone=backbone, neck=neck, head=None) alg = BYOL(backbone=backbone, neck=neck, head=head) - fake_input = torch.randn((16, 3, 224, 224)) + fake_input = torch.randn((2, 3, 224, 224)) fake_backbone_out = alg.extract_feat(fake_input) - assert fake_backbone_out[0].size() == torch.Size([16, 2048, 7, 7]) + assert fake_backbone_out[0].size() == torch.Size([2, 512, 7, 7]) with pytest.raises(AssertionError): fake_out = alg.forward_train(fake_input) - fake_input = [ - torch.randn((16, 3, 224, 224)), - torch.randn((16, 3, 224, 224)) - ] + fake_input = [torch.randn((2, 3, 224, 224)), torch.randn((2, 3, 224, 224))] fake_out = alg.forward_train(fake_input) assert fake_out['loss'].item() > -4 diff --git a/tests/test_models/test_algorithms/test_classification.py b/tests/test_models/test_algorithms/test_classification.py index 132e1a79a..48a5f5d11 100644 --- a/tests/test_models/test_algorithms/test_classification.py +++ b/tests/test_models/test_algorithms/test_classification.py @@ -13,23 +13,23 @@ def test_classification(): with_sobel = True, backbone = dict( type='ResNet', - depth=50, + depth=18, in_channels=2, out_indices=[4], # 0: conv-1, x: stage-x norm_cfg=dict(type='BN'), frozen_stages=4) head = dict( - type='ClsHead', with_avg_pool=True, in_channels=2048, num_classes=4) + type='ClsHead', with_avg_pool=True, in_channels=512, num_classes=4) alg = Classification(backbone=backbone, with_sobel=with_sobel, head=head) assert hasattr(alg, 'sobel_layer') assert hasattr(alg, 'head') - fake_input = torch.randn((16, 3, 224, 224)) - fake_labels = torch.ones(16, dtype=torch.long) + fake_input = torch.randn((2, 3, 224, 224)) + fake_labels = torch.ones(2, dtype=torch.long) fake_out = alg.forward_test(fake_input) assert 'head4' in fake_out - assert fake_out['head4'].size() == torch.Size([16, 4]) + assert fake_out['head4'].size() == torch.Size([2, 4]) fake_out = alg.forward_train(fake_input, fake_labels) assert fake_out['loss'].item() > 0 @@ -51,7 +51,7 @@ def test_classification(): alg = Classification(backbone=backbone, head=head) assert alg.with_head - fake_input = torch.randn((16, 3, 224, 224)) - fake_labels = torch.ones(16, dtype=torch.long) + fake_input = torch.randn((2, 3, 224, 224)) + fake_labels = torch.ones(2, dtype=torch.long) fake_out = alg.forward_train(fake_input, fake_labels) assert fake_out['loss'].item() > 0 diff --git a/tests/test_models/test_algorithms/test_deepcluster.py b/tests/test_models/test_algorithms/test_deepcluster.py index 286fb4c83..f92ee68fc 100644 --- a/tests/test_models/test_algorithms/test_deepcluster.py +++ b/tests/test_models/test_algorithms/test_deepcluster.py @@ -10,7 +10,7 @@ with_sobel = True, backbone = dict( type='ResNet', - depth=50, + depth=18, in_channels=2, out_indices=[4], # 0: conv-1, x: stage-x norm_cfg=dict(type='BN')) @@ -18,7 +18,7 @@ head = dict( type='ClsHead', with_avg_pool=False, # already has avgpool in the neck - in_channels=2048, + in_channels=512, num_classes=num_classes) @@ -34,11 +34,11 @@ def test_deepcluster(): assert hasattr(alg, 'neck') assert hasattr(alg, 'head') - fake_input = torch.randn((16, 3, 224, 224)) - fake_labels = torch.ones(16, dtype=torch.long) + fake_input = torch.randn((2, 3, 224, 224)) + fake_labels = torch.ones(2, dtype=torch.long) fake_out = alg.forward(fake_input, mode='test') assert 'head0' in fake_out - assert fake_out['head0'].size() == torch.Size([16, num_classes]) + assert fake_out['head0'].size() == torch.Size([2, num_classes]) fake_out = alg.forward_train(fake_input, fake_labels) alg.set_reweight(fake_labels) diff --git a/tests/test_models/test_algorithms/test_densecl.py b/tests/test_models/test_algorithms/test_densecl.py index 67820a522..0104a7560 100644 --- a/tests/test_models/test_algorithms/test_densecl.py +++ b/tests/test_models/test_algorithms/test_densecl.py @@ -9,20 +9,20 @@ from mmselfsup.models.algorithms import DenseCL queue_len = 32 -feat_dim = 4 +feat_dim = 2 momentum = 0.999 loss_lambda = 0.5 backbone = dict( type='ResNet', - depth=50, + depth=18, in_channels=3, out_indices=[4], # 0: conv-1, x: stage-x norm_cfg=dict(type='BN')) neck = dict( type='DenseCLNeck', - in_channels=2048, - hid_channels=4, - out_channels=4, + in_channels=512, + hid_channels=2, + out_channels=2, num_grid=None) head = dict(type='ContrastiveHead', temperature=0.2) @@ -57,14 +57,14 @@ def test_densecl(): assert alg.queue.size() == torch.Size([feat_dim, queue_len]) assert alg.queue2.size() == torch.Size([feat_dim, queue_len]) - fake_input = torch.randn((16, 3, 224, 224)) + fake_input = torch.randn((2, 3, 224, 224)) with pytest.raises(AssertionError): fake_out = alg.forward_train(fake_input) fake_out = alg.forward_test(fake_input) assert fake_out[0] is None assert fake_out[2] is None - assert fake_out[1].size() == torch.Size([16, 2048, 49]) + assert fake_out[1].size() == torch.Size([2, 512, 49]) mmselfsup.models.algorithms.densecl.batch_shuffle_ddp = MagicMock( side_effect=mock_batch_shuffle_ddp) @@ -75,10 +75,10 @@ def test_densecl(): fake_loss = alg.forward_train([fake_input, fake_input]) assert fake_loss['loss_single'] > 0 assert fake_loss['loss_dense'] > 0 - assert alg.queue_ptr.item() == 16 - assert alg.queue2_ptr.item() == 16 + assert alg.queue_ptr.item() == 2 + assert alg.queue2_ptr.item() == 2 # 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 + assert fake_outputs['num_samples'] == 2 diff --git a/tests/test_models/test_algorithms/test_mae.py b/tests/test_models/test_algorithms/test_mae.py index d985f44f9..87ad10348 100644 --- a/tests/test_models/test_algorithms/test_mae.py +++ b/tests/test_models/test_algorithms/test_mae.py @@ -30,8 +30,8 @@ def test_mae(): alg = MAE(backbone=None, neck=neck, head=head) alg = MAE(backbone=backbone, neck=neck, head=head) - fake_input = torch.randn((16, 3, 224, 224)) + fake_input = torch.randn((2, 3, 224, 224)) fake_loss = alg.forward_train(fake_input) fake_feature = alg.extract_feat(fake_input) assert isinstance(fake_loss['loss'].item(), float) - assert list(fake_feature[0].shape) == [16, 50, 768] + assert list(fake_feature[0].shape) == [2, 50, 768] diff --git a/tests/test_models/test_algorithms/test_mmcls_classifier_wrapper.py b/tests/test_models/test_algorithms/test_mmcls_classifier_wrapper.py index 114dc3b60..4cd570db9 100644 --- a/tests/test_models/test_algorithms/test_mmcls_classifier_wrapper.py +++ b/tests/test_models/test_algorithms/test_mmcls_classifier_wrapper.py @@ -21,7 +21,7 @@ def test_mmcls_classifier_wrapper(): neck=dict(type='mmcls.GlobalAveragePooling'), head=dict( type='mmcls.LinearClsHead', - num_classes=1000, + num_classes=2, in_channels=1024, init_cfg=None, # suppress the default init_cfg of LinearClsHead. loss=dict( @@ -34,8 +34,8 @@ def test_mmcls_classifier_wrapper(): dict(type='Constant', layer='LayerNorm', val=1., bias=0.) ], train_cfg=dict(augments=[ - dict(type='BatchMixup', alpha=0.8, num_classes=1000, prob=0.5), - dict(type='BatchCutMix', alpha=1.0, num_classes=1000, prob=0.5) + dict(type='BatchMixup', alpha=0.8, num_classes=2, prob=0.5), + dict(type='BatchCutMix', alpha=1.0, num_classes=2, prob=0.5) ])) model = ALGORITHMS.build(model_config) fake_inputs = torch.rand((2, 3, 192, 192)) @@ -47,7 +47,7 @@ def test_mmcls_classifier_wrapper(): # test mode outputs = model(fake_inputs, mode='test') - assert list(outputs['head3'].shape) == [2, 1000] + assert list(outputs['head3'].shape) == [2, 2] # extract mode outputs = model(fake_inputs, mode='extract') diff --git a/tests/test_models/test_algorithms/test_moco.py b/tests/test_models/test_algorithms/test_moco.py index 7735fd3cf..ed29beecf 100644 --- a/tests/test_models/test_algorithms/test_moco.py +++ b/tests/test_models/test_algorithms/test_moco.py @@ -9,19 +9,19 @@ from mmselfsup.models.algorithms import MoCo queue_len = 32 -feat_dim = 4 +feat_dim = 2 momentum = 0.999 backbone = dict( type='ResNet', - depth=50, + depth=18, in_channels=3, out_indices=[4], # 0: conv-1, x: stage-x norm_cfg=dict(type='BN')) neck = dict( type='MoCoV2Neck', - in_channels=2048, - hid_channels=4, - out_channels=4, + in_channels=512, + hid_channels=2, + out_channels=2, with_avg_pool=True) head = dict(type='ContrastiveHead', temperature=0.2) @@ -54,9 +54,9 @@ def test_moco(): momentum=momentum) assert alg.queue.size() == torch.Size([feat_dim, queue_len]) - fake_input = torch.randn((16, 3, 224, 224)) + fake_input = torch.randn((2, 3, 224, 224)) fake_backbone_out = alg.extract_feat(fake_input) - assert fake_backbone_out[0].size() == torch.Size([16, 2048, 7, 7]) + assert fake_backbone_out[0].size() == torch.Size([2, 512, 7, 7]) with pytest.raises(AssertionError): fake_backbone_out = alg.forward_train(fake_input) @@ -68,4 +68,4 @@ def test_moco(): side_effect=mock_concat_all_gather) fake_loss = alg.forward_train([fake_input, fake_input]) assert fake_loss['loss'] > 0 - assert alg.queue_ptr.item() == 16 + assert alg.queue_ptr.item() == 2 diff --git a/tests/test_models/test_algorithms/test_mocov3.py b/tests/test_models/test_algorithms/test_mocov3.py index 74b1842fa..ea42c4fa0 100644 --- a/tests/test_models/test_algorithms/test_mocov3.py +++ b/tests/test_models/test_algorithms/test_mocov3.py @@ -15,8 +15,8 @@ neck = dict( type='NonLinearNeck', in_channels=384, - hid_channels=8, - out_channels=8, + hid_channels=2, + out_channels=2, num_layers=2, with_bias=False, with_last_bn=True, @@ -28,9 +28,9 @@ type='MoCoV3Head', predictor=dict( type='NonLinearNeck', - in_channels=8, - hid_channels=8, - out_channels=8, + in_channels=2, + hid_channels=2, + out_channels=2, num_layers=2, with_bias=False, with_last_bn=True, @@ -51,7 +51,7 @@ def test_mocov3(): alg.init_weights() alg.momentum_update() - fake_input = torch.randn((16, 3, 224, 224)) + fake_input = torch.randn((2, 3, 224, 224)) 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]) + assert fake_backbone_out[0][0].size() == torch.Size([2, 384, 14, 14]) + assert fake_backbone_out[0][1].size() == torch.Size([2, 384]) diff --git a/tests/test_models/test_algorithms/test_npid.py b/tests/test_models/test_algorithms/test_npid.py index 5779903f8..734f58143 100644 --- a/tests/test_models/test_algorithms/test_npid.py +++ b/tests/test_models/test_algorithms/test_npid.py @@ -8,14 +8,14 @@ backbone = dict( type='ResNet', - depth=50, + depth=18, in_channels=3, out_indices=[4], # 0: conv-1, x: stage-x norm_cfg=dict(type='BN')) neck = dict( - type='LinearNeck', in_channels=2048, out_channels=4, with_avg_pool=True) + type='LinearNeck', in_channels=512, out_channels=2, with_avg_pool=True) head = dict(type='ContrastiveHead', temperature=0.07) -memory_bank = dict(type='SimpleMemory', length=8, feat_dim=4, momentum=0.5) +memory_bank = dict(type='SimpleMemory', length=8, feat_dim=2, momentum=0.5) @pytest.mark.skipif( @@ -30,6 +30,6 @@ def test_npid(): alg = NPID( backbone=backbone, neck=neck, head=head, memory_bank=memory_bank) - fake_input = torch.randn((16, 3, 224, 224)) + fake_input = torch.randn((2, 3, 224, 224)) fake_backbone_out = alg.extract_feat(fake_input) - assert fake_backbone_out[0].size() == torch.Size([16, 2048, 7, 7]) + assert fake_backbone_out[0].size() == torch.Size([2, 512, 7, 7]) diff --git a/tests/test_models/test_algorithms/test_odc.py b/tests/test_models/test_algorithms/test_odc.py index b397be857..74d1b7ed8 100644 --- a/tests/test_models/test_algorithms/test_odc.py +++ b/tests/test_models/test_algorithms/test_odc.py @@ -9,26 +9,26 @@ num_classes = 5 backbone = dict( type='ResNet', - depth=50, + depth=18, in_channels=3, out_indices=[4], # 0: conv-1, x: stage-x norm_cfg=dict(type='BN')) neck = dict( type='ODCNeck', - in_channels=2048, - hid_channels=4, - out_channels=4, + in_channels=512, + hid_channels=2, + out_channels=2, norm_cfg=dict(type='BN1d'), with_avg_pool=True) head = dict( type='ClsHead', with_avg_pool=False, - in_channels=4, + in_channels=2, num_classes=num_classes) memory_bank = dict( type='ODCMemory', length=8, - feat_dim=4, + feat_dim=2, momentum=0.5, num_classes=num_classes, min_cluster=2, @@ -48,7 +48,7 @@ def test_odc(): alg = ODC(backbone=backbone, neck=neck, head=head, memory_bank=memory_bank) alg.set_reweight() - fake_input = torch.randn((16, 3, 224, 224)) + fake_input = torch.randn((2, 3, 224, 224)) fake_out = alg.forward_test(fake_input) assert 'head0' in fake_out - assert fake_out['head0'].size() == torch.Size([16, num_classes]) + assert fake_out['head0'].size() == torch.Size([2, num_classes]) diff --git a/tests/test_models/test_algorithms/test_relative_loc.py b/tests/test_models/test_algorithms/test_relative_loc.py index b8feac8f5..1a5f91bce 100644 --- a/tests/test_models/test_algorithms/test_relative_loc.py +++ b/tests/test_models/test_algorithms/test_relative_loc.py @@ -8,16 +8,16 @@ backbone = dict( type='ResNet', - depth=50, + depth=18, in_channels=3, out_indices=[4], # 0: conv-1, x: stage-x norm_cfg=dict(type='BN')) neck = dict( type='RelativeLocNeck', - in_channels=2048, - out_channels=4, + in_channels=512, + out_channels=2, with_avg_pool=True) -head = dict(type='ClsHead', with_avg_pool=False, in_channels=4, num_classes=8) +head = dict(type='ClsHead', with_avg_pool=False, in_channels=2, num_classes=8) @pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit') @@ -49,6 +49,6 @@ def test_relative_loc(): assert 'head4' in fake_out # extract - fake_input = torch.randn((16, 3, 224, 224)) + fake_input = torch.randn((2, 3, 224, 224)) fake_backbone_out = alg.forward(fake_input, mode='extract') - assert fake_backbone_out[0].size() == torch.Size([16, 2048, 7, 7]) + assert fake_backbone_out[0].size() == torch.Size([2, 512, 7, 7]) diff --git a/tests/test_models/test_algorithms/test_rotation_pred.py b/tests/test_models/test_algorithms/test_rotation_pred.py index 69f4b9d3e..46d93a119 100644 --- a/tests/test_models/test_algorithms/test_rotation_pred.py +++ b/tests/test_models/test_algorithms/test_rotation_pred.py @@ -8,12 +8,11 @@ backbone = dict( type='ResNet', - depth=50, + depth=18, in_channels=3, out_indices=[4], # 0: conv-1, x: stage-x norm_cfg=dict(type='BN')) -head = dict( - type='ClsHead', with_avg_pool=True, in_channels=2048, num_classes=4) +head = dict(type='ClsHead', with_avg_pool=True, in_channels=512, num_classes=4) @pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit') @@ -41,6 +40,6 @@ def test_rotation_pred(): assert 'head4' in fake_out # extract - fake_input = torch.randn((16, 3, 224, 224)) + fake_input = torch.randn((2, 3, 224, 224)) fake_backbone_out = alg.forward(fake_input, mode='extract') - assert fake_backbone_out[0].size() == torch.Size([16, 2048, 7, 7]) + assert fake_backbone_out[0].size() == torch.Size([2, 512, 7, 7]) diff --git a/tests/test_models/test_algorithms/test_simclr.py b/tests/test_models/test_algorithms/test_simclr.py index 44744aa26..3ccc27d81 100644 --- a/tests/test_models/test_algorithms/test_simclr.py +++ b/tests/test_models/test_algorithms/test_simclr.py @@ -8,15 +8,15 @@ backbone = dict( type='ResNet', - depth=50, + depth=18, in_channels=3, out_indices=[4], # 0: conv-1, x: stage-x norm_cfg=dict(type='BN')) neck = dict( type='NonLinearNeck', # SimCLR non-linear neck - in_channels=2048, - hid_channels=4, - out_channels=4, + in_channels=512, + hid_channels=2, + out_channels=2, num_layers=2, with_avg_pool=True) head = dict(type='ContrastiveHead', temperature=0.1) @@ -34,6 +34,6 @@ def test_simclr(): fake_input = torch.randn((16, 3, 224, 224)) alg.forward_train(fake_input) - fake_input = torch.randn((16, 3, 224, 224)) + fake_input = torch.randn((2, 3, 224, 224)) fake_backbone_out = alg.extract_feat(fake_input) - assert fake_backbone_out[0].size() == torch.Size([16, 2048, 7, 7]) + assert fake_backbone_out[0].size() == torch.Size([2, 512, 7, 7]) diff --git a/tests/test_models/test_algorithms/test_simsiam.py b/tests/test_models/test_algorithms/test_simsiam.py index fe2011dfd..93469fbc6 100644 --- a/tests/test_models/test_algorithms/test_simsiam.py +++ b/tests/test_models/test_algorithms/test_simsiam.py @@ -8,16 +8,16 @@ backbone = dict( type='ResNet', - depth=50, + depth=18, in_channels=3, out_indices=[4], # 0: conv-1, x: stage-x norm_cfg=dict(type='BN'), zero_init_residual=True) neck = dict( type='NonLinearNeck', - in_channels=2048, - hid_channels=4, - out_channels=4, + in_channels=512, + hid_channels=2, + out_channels=2, num_layers=3, with_last_bn_affine=False, with_avg_pool=True, @@ -26,9 +26,9 @@ type='LatentPredictHead', predictor=dict( type='NonLinearNeck', - in_channels=4, - hid_channels=4, - out_channels=4, + in_channels=2, + hid_channels=2, + out_channels=2, with_avg_pool=False, with_last_bn=False, with_last_bias=True, @@ -42,22 +42,19 @@ def test_simsiam(): alg = SimSiam(backbone=backbone, neck=neck, head=head) with pytest.raises(AssertionError): - fake_input = torch.randn((16, 3, 224, 224)) + fake_input = torch.randn((2, 3, 224, 224)) alg.forward_train(fake_input) - fake_input = [ - torch.randn((16, 3, 224, 224)), - torch.randn((16, 3, 224, 224)) - ] + fake_input = [torch.randn((2, 3, 224, 224)), torch.randn((2, 3, 224, 224))] 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 + assert fake_outputs['num_samples'] == 2 # 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 + assert fake_outputs['num_samples'] == 2 diff --git a/tests/test_models/test_algorithms/test_swav.py b/tests/test_models/test_algorithms/test_swav.py index 0d09db5d8..09a80809f 100644 --- a/tests/test_models/test_algorithms/test_swav.py +++ b/tests/test_models/test_algorithms/test_swav.py @@ -9,21 +9,21 @@ nmb_crops = [2, 6] backbone = dict( type='ResNet', - depth=50, + depth=18, in_channels=3, out_indices=[4], # 0: conv-1, x: stage-x norm_cfg=dict(type='BN'), zero_init_residual=True) neck = dict( type='SwAVNeck', - in_channels=2048, - hid_channels=4, - out_channels=4, + in_channels=512, + hid_channels=2, + out_channels=2, norm_cfg=dict(type='BN1d'), with_avg_pool=True) head = dict( type='SwAVHead', - feat_dim=4, # equal to neck['out_channels'] + feat_dim=2, # equal to neck['out_channels'] epsilon=0.05, temperature=0.1, num_crops=nmb_crops) @@ -37,19 +37,19 @@ def test_swav(): alg = SwAV(backbone=backbone, neck=None, head=head) alg = SwAV(backbone=backbone, neck=neck, head=head) - fake_input = torch.randn((16, 3, 224, 224)) + fake_input = torch.randn((2, 3, 224, 224)) fake_backbone_out = alg.extract_feat(fake_input) - assert fake_backbone_out[0].size() == torch.Size([16, 2048, 7, 7]) + assert fake_backbone_out[0].size() == torch.Size([2, 512, 7, 7]) fake_input = [ - torch.randn((16, 3, 224, 224)), - torch.randn((16, 3, 224, 224)), - torch.randn((16, 3, 96, 96)), - torch.randn((16, 3, 96, 96)), - torch.randn((16, 3, 96, 96)), - torch.randn((16, 3, 96, 96)), - torch.randn((16, 3, 96, 96)), - torch.randn((16, 3, 96, 96)), + torch.randn((2, 3, 224, 224)), + torch.randn((2, 3, 224, 224)), + torch.randn((2, 3, 96, 96)), + torch.randn((2, 3, 96, 96)), + torch.randn((2, 3, 96, 96)), + torch.randn((2, 3, 96, 96)), + torch.randn((2, 3, 96, 96)), + torch.randn((2, 3, 96, 96)), ] fake_out = alg.forward_train(fake_input) assert fake_out['loss'].item() > 0