My loss seems to be not decreasing for FCN16 and FCN8, it gets stuck at a constant loss value. Here is the code snippet of my Model. Please let know if the model is correct for FCN16 / FCN8. I am trying to do Image segmentation on biomedical images. I believe my model is correct but still would like to have an additional pair of eyes to verify it. I would also need some suggestion on pretrained features, currently I am not taking any pretrained features.
class FCN(nn.Module):
def init(self, num_classes , mode=‘fcn32’):
super(FCN,self).init()
self.mode = mode
feats = list(models.vgg16().features.children())
if self.mode == 'fcn16':
self.feats = nn.Sequential(*feats[0:17])
self.pool4 = nn.Sequential(*feats[17:24])
self.pool5 = nn.Sequential(*feats[24:])
self.res_pool4 = nn.Conv2d(512, num_classes, 1)
elif self.mode == 'fcn8':
self.feats = nn.Sequential(*feats[0:10])
self.pool3 = nn.Sequential(*feats[10:17])
self.pool4 = nn.Sequential(*feats[17:24])
self.pool5 = nn.Sequential(*feats[24:])
self.res_pool3 = nn.Conv2d(256, num_classes, 1)
self.res_pool4 = nn.Conv2d(512, num_classes, 1)
else :
self.feats = nn.Sequential(*feats)
self.fconn = nn.Sequential(nn.Conv2d(512, 4096, 7, padding=3),
nn.ReLU(inplace=True),
nn.Conv2d(4096, 4096,1),
nn.ReLU(inplace=True),
nn.Conv2d(4096, num_classes, 1)
)
self.activation = nn.Sigmoid()
def forward(self,x):
feats = self.feats(x)
if self.mode == 'fcn8':
pool3 = self.pool3(feats)
pool4 = self.pool4(pool3)
pool5 = self.pool5(pool4)
fconn = self.fconn(pool5)
res_pool3 = self.res_pool3(pool3)
res_pool4 = self.res_pool4(pool4)
# Upsampling to match the shape of 3rd Pooling Layer
up_pool4 = func.upsample(res_pool4, pool3.size()[2:], mode='bilinear')
# Upsampling to match the shape of 3rd Pooling Layer
up_fconn = func.upsample(fconn, pool3.size()[2:], mode='bilinear')
out = up_pool4 + up_fconn + res_pool3
elif self.mode == 'fcn16':
pool4 = self.pool4(feats)
pool5 = self.pool5(pool4)
fconn = self.fconn(pool5)
res_pool4 = self.res_pool4(pool4)
# Upsampling to match the shape of 4th Pooling Layer
up_fconn = func.upsample(fconn, pool4.size()[2:], mode='bilinear')
out = up_fconn + res_pool4
else:
out = self.fconn(feats)
upsample = func.upsample(out, x.size()[2:],mode='bilinear')
out = self.activation(upsample)
return out