mirror of
https://github.com/open-mmlab/mmselfsup.git
synced 2025-06-03 14:59:38 +08:00
6732025f48
* DenseCL init weights copy query encoder weights to key encoder. * Logger prints that key encoder is initialized with query encoder.
91 lines
2.6 KiB
Python
91 lines
2.6 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import platform
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import mmselfsup
|
|
from mmselfsup.models.algorithms import DenseCL
|
|
|
|
queue_len = 32
|
|
feat_dim = 2
|
|
momentum = 0.999
|
|
loss_lambda = 0.5
|
|
backbone = dict(
|
|
type='ResNet',
|
|
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=512,
|
|
hid_channels=2,
|
|
out_channels=2,
|
|
num_grid=None)
|
|
head = dict(type='ContrastiveHead', temperature=0.2)
|
|
|
|
|
|
def mock_batch_shuffle_ddp(img):
|
|
return img, 0
|
|
|
|
|
|
def mock_batch_unshuffle_ddp(img, mock_input):
|
|
return img
|
|
|
|
|
|
def mock_concat_all_gather(img):
|
|
return img
|
|
|
|
|
|
@pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
|
|
def test_densecl():
|
|
with pytest.raises(AssertionError):
|
|
alg = DenseCL(backbone=backbone, neck=None, head=head)
|
|
with pytest.raises(AssertionError):
|
|
alg = DenseCL(backbone=backbone, neck=neck, head=None)
|
|
|
|
alg = DenseCL(
|
|
backbone=backbone,
|
|
neck=neck,
|
|
head=head,
|
|
queue_len=queue_len,
|
|
feat_dim=feat_dim,
|
|
momentum=momentum,
|
|
loss_lambda=loss_lambda)
|
|
assert alg.queue.size() == torch.Size([feat_dim, queue_len])
|
|
assert alg.queue2.size() == torch.Size([feat_dim, queue_len])
|
|
|
|
alg.init_weights()
|
|
for param_q, param_k in zip(alg.encoder_q.parameters(),
|
|
alg.encoder_k.parameters()):
|
|
assert torch.equal(param_q, param_k)
|
|
assert param_k.requires_grad is False
|
|
|
|
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([2, 512, 49])
|
|
|
|
mmselfsup.models.algorithms.densecl.batch_shuffle_ddp = MagicMock(
|
|
side_effect=mock_batch_shuffle_ddp)
|
|
mmselfsup.models.algorithms.densecl.batch_unshuffle_ddp = MagicMock(
|
|
side_effect=mock_batch_unshuffle_ddp)
|
|
mmselfsup.models.algorithms.densecl.concat_all_gather = MagicMock(
|
|
side_effect=mock_concat_all_gather)
|
|
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() == 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'] == 2
|