Warmup epochs with frozen pre-trained encoder weights to initialize decoder

[daniel-j-h]

At the moment we are using a pre-trained ResNet as an encoder in our encoder-decoder architecture:

robosat/robosat/unet.py

Lines 94 to 100 in 8b7566e

	self.resnet = resnet50(pretrained=pretrained)
	self.enc0 = nn.Sequential(self.resnet.conv1, self.resnet.bn1, self.resnet.relu, self.resnet.maxpool)
	self.enc1 = self.resnet.layer1 # 256
	self.enc2 = self.resnet.layer2 # 512
	self.enc3 = self.resnet.layer3 # 1024
	self.enc4 = self.resnet.layer4 # 2048

robosat/robosat/unet.py

Lines 123 to 134 in 8b7566e

	enc0 = self.enc0(x)
	enc1 = self.enc1(enc0)
	enc2 = self.enc2(enc1)
	enc3 = self.enc3(enc2)
	enc4 = self.enc4(enc3)
	center = self.center(nn.functional.max_pool2d(enc4, kernel_size=2, stride=2))
	dec0 = self.dec0(torch.cat([enc4, center], dim=1))
	dec1 = self.dec1(torch.cat([enc3, dec0], dim=1))
	dec2 = self.dec2(torch.cat([enc2, dec1], dim=1))
	dec3 = self.dec3(torch.cat([enc1, dec2], dim=1))

We are currently training the model as is with all layers unfrozen.

We should investigate if freezing the ResNet encoder and running a few warmup epochs to initialize the decoder layers (then unfreezing parts or all of the ResNet) helps.

Here is how we can freeze the encoder - unfreezing works similarly:

def freeze(self):
    for layer in (self.enc0, self.enc1, self.enc2, self.enc3, self.enc4):
        for param in layer.parameters():
            param.requires_grad = False