Rearranging pretrained network

I am trying to pretrain a neural network on say cifar before using this trained backbone on another task.

This is trivial if I just need to chop off the last layer and add new layers but how can I approach this to add in skip connections to new layers?

import torch.nn as nn
import torch.nn.functional as F

class ExampleBackbone(nn.Module):
    def __init__(self, num_classes=100):
          self.conv1 = nn.Conv2d(3, 16, kernel_size=3,stride=1,padding=1)
          self.bn1 = nn.BatchNorm2d(ch_out, momentum=0.01)
          self.maxpool = nn.MaxPool2d(2,2)
          self.linear = nn.Linear(16, 100)

    def forward(self, x):
          out = self.conv1(x)
          out = self.bn1(out)
          out = self.maxpool(out)
          return F.log_softmax(self.linear(out))

So for my final result I might want to take the output from self.bn1 and plug to a new layer that I have added for my specific task. Is that the best way to structure the backbone?


If you add other layers to your model, you would have to set strict=False while loading the state_dict to avoid missing keys errors.
The skip connections might just be added in your forward method, which should work fine.

so do I create a new child class to my original or?

Yeah, sounds like a good idea!
Let me know, if you get stuck somewhere.

the self.bn1 though isn’t saving the result anywhere though? guessing I’ll have to save the result to some class variable somewhere to use it? or do I use a pretrained_model.__getitem(x)__ type method in the new layer?

I’m not sure to understand your question correctly.
If you would like to store the output of self.bn1 to use it as a skip connection, you can just use another variable name:

    def forward(self, x):
          out = self.conv1(x)
          bn_out = self.bn1(out)
          out = self.maxpool(bn_out)
          return F.log_softmax(self.linear(out))

so my question is I’ve done that and trained a network using say CIFAR-10 using this network now I want to take this network and add yolo layers that will use bn_out as a skip connection will that work?

I’m not really familiar with the architecture of YOLO and am nut sure, how exactly you would like to add skip connections between your two models.
Anyway, based on your example model, you could just add layers to your model (e.g. from YOLO) and manipulate the forward method as you wish.