Adding an auxiliary branch as a regulariser

vision
1

I am trying to implement the description from this paper: Regularizing Deep Networks by Modeling and Predicting Label Structure. I have designed my autoencoder and FCN, but I am struggling with trying to combine them. I guess I want to load in my pretrained autoencoder and somehow only run the decoder half of the network (as specified in Figure 1 of the paper). Is it possible to do something like this in PyTorch with a pretrained model:

class AutoEncoderConv2(nn.Module):
    def __init__(self):
        super(AutoEncoderConv, self).__init__()

        self.encoder = nn.Sequential(

            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.MaxPool2d(2),

            nn.Conv2d(32, 32, kernel_size=3, padding=1),
            nn.MaxPool2d(2),

            nn.Conv2d(32, 32, kernel_size=3, padding=1),
            nn.MaxPool2d(2),

            nn.Conv2d(32, 32, kernel_size=3, padding=1),
            nn.MaxPool2d(2),

            nn.Conv2d(32, 32, kernel_size=3, padding=1),
            nn.MaxPool2d(2)
        )

        self.decoder = nn.Sequential(

            Interpolate(mode='bilinear', scale_factor=2),
            nn.ConvTranspose2d(32, 32, kernel_size=3, padding=1),

            Interpolate(mode='bilinear', scale_factor=2),
            nn.ConvTranspose2d(32, 32, kernel_size=3, padding=1),

            Interpolate(mode='bilinear', scale_factor=2),
            nn.ConvTranspose2d(32, 32, kernel_size=3, padding=1),

            Interpolate(mode='bilinear', scale_factor=2),
            nn.ConvTranspose2d(32, 32, kernel_size=3, padding=1),

            Interpolate(mode='bilinear', scale_factor=2),
            nn.ConvTranspose2d(32, 1, kernel_size=3, padding=1),
            nn.ReLU(True),

            nn.Tanh()
        )

    def encode(self, x):
        return self.encoder(x)

    def decode(self, y):
        return self.decoder(y)

    def forward(self, x):
        mid_x = self.encoder(x)
        final_x = self.decoder(mid_x)

        return final_x

and then in the FCN code, do this:

model2 = AutoEncoderConv2()
model2.load_state_dict(torch.load("trained_models/autoencoder.pth"))

fcn_model = FCN8s(pretrained_net=vgg_model, n_class=2)


#inside training loop
optimizer.zero_grad()
inputs = Variable(inputs.cuda())
labels = Variable(labels.cuda())

outputs = fcn_model(inputs)

#and then to only use the decoder:
outputs2 = model2.decoder(intermediate_input_from_fcn)

#Freeze decoder 
for p in model2.parameters():
      p.requires_grad = False

loss1 = criterion(outputs, labels)
loss2 = criterion(outputs2, labels)
loss = loss1 + loss2

loss.backward()

I need to only use the decoder part of the trained autoencoder and then backpropagate the losses through the FCN (not the decoder), so the decoder parameters should not be updated (i.e. decoder should be frozen).

Any help would be appreciated, thank you.

2

I think you are almost there. Place this code before using model2. i.e., at the top after load_state_dict.

(or)

Do not specify the model2 parameters to the optimizer. i.e., only specify FCN parameters.

1 Like