def patch_torchvision_mobilenet_v2(model):
"""
Patches TorchVision's MobileNetV2:
* To allow quantization, this adds modules for tensor operations (mean, element-wise addition) to the
model instance and patches the forward functions accordingly
* Fixes a bug in the torchvision implementation that prevents export to ONNX (and creation of SummaryGraph)
"""
if not isinstance(model, torch_models.MobileNetV2):
raise TypeError("Only MobileNetV2 is acceptable.")
def patched_forward_mobilenet_v2(self, x):
x = self.features(x)
# x = x.mean([2, 3]) # this was a bug: https://github.com/pytorch/pytorch/issues/20516
x = self.mean32(x)
x = self.classifier(x)
return x
model.mean32 = nn.Sequential(
Mean(3), Mean(2)
)
model.__class__.forward = patched_forward_mobilenet_v2
def is_inverted_residual(module):
return isinstance(module, nn.Module) and module.__class__.__name__ == 'InvertedResidual'
def patched_forward_invertedresidual(self, x):
if self.use_res_connect:
return self.residual_eltwiseadd(self.conv(x), x)
else:
return self.conv(x)
for n, m in model.named_modules():
if is_inverted_residual(m):
if m.use_res_connect:
m.residual_eltwiseadd = EltwiseAdd()
m.__class__.forward = patched_forward_invertedresidual
def __init__(self, *args, **kwargs):
# Initialize torchvision version
super(DistillerBottleneck, self).__init__(*args, **kwargs)
# Remove original relu in favor of numbered modules
delattr(self, 'relu')
self.relu1 = nn.ReLU(inplace=True)
self.relu2 = nn.ReLU(inplace=True)
self.relu3 = nn.ReLU(inplace=True)
self.add = EltwiseAdd(
inplace=True) # Replace '+=' operator with inplace module
# Trick to make the modules accessible in their topological order
modules = OrderedDict()
modules['conv1'] = self.conv1
modules['bn1'] = self.bn1
modules['relu1'] = self.relu1
modules['conv2'] = self.conv2
modules['bn2'] = self.bn2
modules['relu2'] = self.relu2
modules['conv3'] = self.conv3
modules['bn3'] = self.bn3
if self.downsample is not None:
modules['downsample'] = self.downsample
modules['add'] = self.add
modules['relu3'] = self.relu3
self._modules = modules
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(BasicBlock, self).__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn1 = nn.BatchNorm2d(planes)
self.relu1 = nn.ReLU(inplace=True)
self.conv2 = conv3x3(planes, planes)
self.bn2 = nn.BatchNorm2d(planes)
self.relu2 = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
# Replace '+=' operator with inplace module
self.add = EltwiseAdd(inplace=True)
def __init__(self,
block_gates,
inplanes,
planes,
stride=1,
downsample=None):
super(BasicBlock, self).__init__()
self.block_gates = block_gates
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn1 = nn.BatchNorm2d(planes)
self.relu1 = nn.ReLU(
inplace=False) # To enable layer removal inplace must be False
self.conv2 = conv3x3(planes, planes)
self.bn2 = nn.BatchNorm2d(planes)
self.relu2 = nn.ReLU(inplace=False)
self.downsample = downsample
self.stride = stride
self.residual_eltwiseadd = EltwiseAdd()
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(Bottleneck, self).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.relu1 = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(
planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
)
self.bn2 = nn.BatchNorm2d(planes)
self.relu2 = nn.ReLU(inplace=True)
self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
self.bn3 = nn.BatchNorm2d(planes * 4)
self.relu3 = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
# Replace '+=' operator with inplace module
self.add = EltwiseAdd(inplace=True)
def __init__(self, inp, oup, stride, expand_ratio):
super(InvertedResidual, self).__init__()
self.stride = stride
assert stride in [1, 2]
hidden_dim = int(round(inp * expand_ratio))
self.use_res_connect = self.stride == 1 and inp == oup
layers = []
if expand_ratio != 1:
# pw
layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
layers.extend([
# dw
ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
# pw-linear
nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
nn.BatchNorm2d(oup),
])
self.conv = nn.Sequential(*layers)
self.add = EltwiseAdd(inplace=True)
def __init__(self, prologue, bypassed):
super(BypassModel, self).__init__()
self.prologue = nn.Sequential(*prologue)
self.bypassed = bypassed
self.add = EltwiseAdd()
def __init__(self, m1, m2):
super(SplitJoinModel, self).__init__()
self.split = Split(int(input_shape[0] / 2))
self.m1 = m1
self.m2 = m2
self.add = EltwiseAdd()