How can I add a module into a network without change the main body of the network?

user_13316883 0
Viewed 171

How can I add a module into a network without change the main body of the network? For example, I want to add an attention module after the layer1 of Resnet. I have tried to use forward_hook to add the new module, but after reading the source code of nn.Module: 

def __call__(self, *input, **kwargs):
   result = self.forward(*input, **kwargs)
   for hook in self._forward_hooks.values():
       hook_result = hook(self, input, result)
   ...
   return result

I found that hooks are done after forward, so I can't use hook to add new modules. Are there any other way to realize this function?

pytorch python
1 Answers

The key point is make a new forward function and ditch the old one. You can inherit the ResNet directly and just add a new layer. Then specific how the forward should be.

Like

class AttentionResNet(torchvision.models.resnet.ResNet):
    def __init__(self, *args, **kwargs):
        super(AttentionResNet, self).__init__(*args, **kwargs)
        self.attention = Attention()

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        
        # this is what you asked
        x = self.attention(x)
        
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x

Or if you want pretrained weights you can also make the resnet as a sub module and don't call the resnet's forward directly, like.

class AttentionResNet(nn.Module):
    def __init__(self):
        super(AttentionResNet, self).__init__()
        self.attention = Attention()
        self.resnet = torchvision.models.resnet.resnet50()

    def forward(self, x):
        x = self.resnet.conv1(x)
        x = self.resnet.bn1(x)
        x = self.resnet.relu(x)
        x = self.resnet.maxpool(x)

        x = self.resnet.layer1(x)

        # this is what you want
        x = self.attention(x)

        x = self.resnet.layer2(x)
        x = self.resnet.layer3(x)
        x = self.resnet.layer4(x)

        x = self.resnet.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.resnet.fc(x)
        return x

Note that the forward in ResNet looks like this. And for me using forward directly doesn't cause problem but you can also use the _forward_impl() like original if you want to.

    def _forward_impl(self, x):
        # See note [TorchScript super()]
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x

    def forward(self, x):
        return self._forward_impl(x)
user_94773380
Related