Spaces:
Sleeping
Sleeping
File size: 996 Bytes
d7e58f0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
# Copyright (c) OpenMMLab. All rights reserved.
from unittest.mock import MagicMock, patch
import torch.nn as nn
from mmcv.device.mlu import MLUDataParallel, MLUDistributedDataParallel
from mmcv.parallel import is_module_wrapper
from mmcv.utils import IS_MLU_AVAILABLE
def mock(*args, **kwargs):
pass
@patch('torch.distributed._broadcast_coalesced', mock)
@patch('torch.distributed.broadcast', mock)
@patch('torch.nn.parallel.DistributedDataParallel._ddp_init_helper', mock)
def test_is_module_wrapper():
class Model(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(2, 2, 1)
def forward(self, x):
return self.conv(x)
model = Model()
assert not is_module_wrapper(model)
if IS_MLU_AVAILABLE:
mludp = MLUDataParallel(model)
assert is_module_wrapper(mludp)
mluddp = MLUDistributedDataParallel(model, process_group=MagicMock())
assert is_module_wrapper(mluddp)
|