25 lines
739 B
Python
25 lines
739 B
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from mmselfsup.models.necks import MoCoV2Neck
|
|
|
|
|
|
def test_mocov2_neck():
|
|
neck = MoCoV2Neck(16, 32, 16)
|
|
assert isinstance(neck.mlp, nn.Sequential)
|
|
assert neck.mlp[0].in_features == 16
|
|
assert neck.mlp[2].in_features == 32
|
|
assert neck.mlp[2].out_features == 16
|
|
|
|
# test neck with avgpool
|
|
fake_in = torch.rand((32, 16, 5, 5))
|
|
fake_out = neck.forward([fake_in])
|
|
assert fake_out[0].shape == torch.Size([32, 16])
|
|
|
|
# test neck without avgpool
|
|
neck = MoCoV2Neck(16, 32, 16, with_avg_pool=False)
|
|
fake_in = torch.rand((32, 16))
|
|
fake_out = neck.forward([fake_in])
|
|
assert fake_out[0].shape == torch.Size([32, 16])
|