def test_replace_basic_case(self):
model = nn.BatchNorm1d(10)
model = mm.replace_all_modules(
model, nn.BatchNorm1d, lambda _: nn.BatchNorm2d(10)
)
self.checkModulePresent(model, nn.BatchNorm2d)
self.checkModuleNotPresent(model, nn.BatchNorm1d)
def test_replace_sequential_case(self):
model = nn.Sequential(nn.Conv1d(1, 2, 3), nn.Sequential(nn.Conv2d(4, 5, 6)))
def conv(m: nn.Conv2d):
return nn.Linear(4, 5)
model = mm.replace_all_modules(model, nn.Conv2d, conv)
self.checkModulePresent(model, nn.Linear)
self.checkModuleNotPresent(model, nn.Conv2d)
def test_module_modification_replace_example(self):
# IMPORTANT: When changing this code you also need to update
# the docstring for opacus.utils.module_modification.replace_all_modules()
from torchvision.models import resnet18
model = resnet18()
self.assertTrue(isinstance(model.layer1[0].bn1, nn.BatchNorm2d))
model = replace_all_modules(model, nn.BatchNorm2d, lambda _: nn.Identity())
self.assertTrue(isinstance(model.layer1[0].bn1, nn.Identity))