# # from https://github.com/wkentaro/pytorch-fcn import torch.nn as nn from torch.nn import functional as F class FCN8s(nn.Module): def __init__(self, n_classes=4, n_frames=5): super(FCN8s, self).__init__() # lt@20240406 reduce resolution self.pool0 = nn.MaxPool2d(2, stride=2, ceil_mode=True) # conv1 # The 100 pixel input padding guarantees that the network output can be aligned to the input for any input size # since the kernel size of fc6 is 7 and its padding equals to 0. self.conv1_1 = nn.Conv2d(n_frames, 64, 3, padding=100) self.relu1_1 = nn.ReLU(inplace=True) self.conv1_2 = nn.Conv2d(64, 64, 3, padding=1) self.relu1_2 = nn.ReLU(inplace=True) self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True) # 1/2 # conv2 self.conv2_1 = nn.Conv2d(64, 128, 3, padding=1) self.relu2_1 = nn.ReLU(inplace=True) self.conv2_2 = nn.Conv2d(128, 128, 3, padding=1) self.relu2_2 = nn.ReLU(inplace=True) self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True) # 1/4 # conv3 self.conv3_1 = nn.Conv2d(128, 256, 3, padding=1) self.relu3_1 = nn.ReLU(inplace=True) self.conv3_2 = nn.Conv2d(256, 256, 3, padding=1) self.relu3_2 = nn.ReLU(inplace=True) self.conv3_3 = nn.Conv2d(256, 256, 3, padding=1) self.relu3_3 = nn.ReLU(inplace=True) self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True) # 1/8 # conv4 self.conv4_1 = nn.Conv2d(256, 512, 3, padding=1) self.relu4_1 = nn.ReLU(inplace=True) self.conv4_2 = nn.Conv2d(512, 512, 3, padding=1) self.relu4_2 = nn.ReLU(inplace=True) self.conv4_3 = nn.Conv2d(512, 512, 3, padding=1) self.relu4_3 = nn.ReLU(inplace=True) self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True) # 1/16 # conv5 self.conv5_1 = nn.Conv2d(512, 512, 3, padding=1) self.relu5_1 = nn.ReLU(inplace=True) self.conv5_2 = nn.Conv2d(512, 512, 3, padding=1) self.relu5_2 = nn.ReLU(inplace=True) self.conv5_3 = nn.Conv2d(512, 512, 3, padding=1) self.relu5_3 = nn.ReLU(inplace=True) self.pool5 = nn.MaxPool2d(2, stride=2, ceil_mode=True) # 1/32 # fc6 self.fc6 = nn.Conv2d(512, 4096, 7) self.relu6 = nn.ReLU(inplace=True) self.drop6 = nn.Dropout2d() # fc7 self.fc7 = nn.Conv2d(4096, 4096, 1) self.relu7 = nn.ReLU(inplace=True) self.drop7 = nn.Dropout2d() self.score_fr = nn.Conv2d(4096, n_classes, 1) self.score_pool3 = nn.Conv2d(256, n_classes, 1) self.score_pool4 = nn.Conv2d(512, n_classes, 1) self.upscore2 = nn.ConvTranspose2d( n_classes, n_classes, 4, stride=2, bias=False) self.upscore8 = nn.ConvTranspose2d( n_classes, n_classes, 16, stride=8, bias=False) self.upscore_pool4 = nn.ConvTranspose2d( n_classes, n_classes, 4, stride=2, bias=False) def forward(self, x): # lt@20240406 reduce resolution x = self.pool0(x) # lt@20240408 add resolution # x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True) h = x h = self.relu1_1(self.conv1_1(h)) h = self.relu1_2(self.conv1_2(h)) h = self.pool1(h) h = self.relu2_1(self.conv2_1(h)) h = self.relu2_2(self.conv2_2(h)) h = self.pool2(h) h = self.relu3_1(self.conv3_1(h)) h = self.relu3_2(self.conv3_2(h)) h = self.relu3_3(self.conv3_3(h)) h = self.pool3(h) pool3 = h # 1/8 h = self.relu4_1(self.conv4_1(h)) h = self.relu4_2(self.conv4_2(h)) h = self.relu4_3(self.conv4_3(h)) h = self.pool4(h) pool4 = h # 1/16 h = self.relu5_1(self.conv5_1(h)) h = self.relu5_2(self.conv5_2(h)) h = self.relu5_3(self.conv5_3(h)) h = self.pool5(h) h = self.relu6(self.fc6(h)) h = self.drop6(h) h = self.relu7(self.fc7(h)) h = self.drop7(h) h = self.score_fr(h) h = self.upscore2(h) upscore2 = h # 1/16 h = self.score_pool4(pool4) # crop for alignment h = h[:, :, 5:5 + upscore2.size()[2], 5:5 + upscore2.size()[3]] score_pool4c = h # 1/16 h = upscore2 + score_pool4c # 1/16 h = self.upscore_pool4(h) upscore_pool4 = h # 1/8 h = self.score_pool3(pool3) # crop h = h[:, :, 9:9 + upscore_pool4.size()[2], 9:9 + upscore_pool4.size()[3]] score_pool3c = h # 1/8 h = upscore_pool4 + score_pool3c # 1/8 h = self.upscore8(h) # crop h = h[:, :, 31:31 + x.size()[2], 31:31 + x.size()[3]].contiguous() # lt@20240406 recover resolution h = F.interpolate(h, scale_factor=2, mode='bilinear', align_corners=True) # lt@20240408 add resolution # h = F.interpolate(h, scale_factor=0.5, mode='bilinear', align_corners=True) return h # from https://github.com/wkentaro/pytorch-fcn # import torch.nn as nn # class FCN8s(nn.Module): # def __init__(self, n_classes=4, n_frames=5): # super(FCN8s, self).__init__() # # conv1 # # The 100 pixel input padding guarantees that the network output can be aligned to the input for any input size # # since the kernel size of fc6 is 7 and its padding equals to 0. # self.conv1_1 = nn.Conv2d(n_frames, 64, 3, padding=100) # self.bn1_1 = nn.BatchNorm2d(64) # self.relu1_1 = nn.ReLU(inplace=True) # self.conv1_2 = nn.Conv2d(64, 64, 3, padding=1) # self.bn1_2 = nn.BatchNorm2d(64) # self.relu1_2 = nn.ReLU(inplace=True) # self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True) # 1/2 # # conv2 # self.conv2_1 = nn.Conv2d(64, 128, 3, padding=1) # self.bn2_1 = nn.BatchNorm2d(128) # self.relu2_1 = nn.ReLU(inplace=True) # self.conv2_2 = nn.Conv2d(128, 128, 3, padding=1) # self.bn2_2 = nn.BatchNorm2d(128) # self.relu2_2 = nn.ReLU(inplace=True) # self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True) # 1/4 # # conv3 # self.conv3_1 = nn.Conv2d(128, 256, 3, padding=1) # self.bn3_1 = nn.BatchNorm2d(256) # self.relu3_1 = nn.ReLU(inplace=True) # self.conv3_2 = nn.Conv2d(256, 256, 3, padding=1) # self.bn3_2 = nn.BatchNorm2d(256) # self.relu3_2 = nn.ReLU(inplace=True) # self.conv3_3 = nn.Conv2d(256, 256, 3, padding=1) # self.bn3_3 = nn.BatchNorm2d(256) # self.relu3_3 = nn.ReLU(inplace=True) # self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True) # 1/8 # # conv4 # self.conv4_1 = nn.Conv2d(256, 512, 3, padding=1) # self.bn4_1 = nn.BatchNorm2d(512) # self.relu4_1 = nn.ReLU(inplace=True) # self.conv4_2 = nn.Conv2d(512, 512, 3, padding=1) # self.bn4_2 = nn.BatchNorm2d(512) # self.relu4_2 = nn.ReLU(inplace=True) # self.conv4_3 = nn.Conv2d(512, 512, 3, padding=1) # self.bn4_3 = nn.BatchNorm2d(512) # self.relu4_3 = nn.ReLU(inplace=True) # self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True) # 1/16 # # conv5 # self.conv5_1 = nn.Conv2d(512, 512, 3, padding=1) # self.bn5_1 = nn.BatchNorm2d(512) # self.relu5_1 = nn.ReLU(inplace=True) # self.conv5_2 = nn.Conv2d(512, 512, 3, padding=1) # self.bn5_2 = nn.BatchNorm2d(512) # self.relu5_2 = nn.ReLU(inplace=True) # self.conv5_3 = nn.Conv2d(512, 512, 3, padding=1) # self.bn5_3 = nn.BatchNorm2d(512) # self.relu5_3 = nn.ReLU(inplace=True) # self.pool5 = nn.MaxPool2d(2, stride=2, ceil_mode=True) # 1/32 # # fc6 # self.fc6 = nn.Conv2d(512, 4096, 7) # self.bn6 = nn.BatchNorm2d(4096) # self.relu6 = nn.ReLU(inplace=True) # self.drop6 = nn.Dropout2d() # # fc7 # self.fc7 = nn.Conv2d(4096, 4096, 1) # self.bn7 = nn.BatchNorm2d(4096) # self.relu7 = nn.ReLU(inplace=True) # self.drop7 = nn.Dropout2d() # self.score_fr = nn.Conv2d(4096, n_classes, 1) # self.score_pool3 = nn.Conv2d(256, n_classes, 1) # self.score_pool4 = nn.Conv2d(512, n_classes, 1) # self.upscore2 = nn.ConvTranspose2d( # n_classes, n_classes, 4, stride=2, bias=False) # self.upscore8 = nn.ConvTranspose2d( # n_classes, n_classes, 16, stride=8, bias=False) # self.upscore_pool4 = nn.ConvTranspose2d( # n_classes, n_classes, 4, stride=2, bias=False) # def forward(self, x): # h = x # h = self.relu1_1(self.bn1_1(self.conv1_1(h))) # h = self.relu1_2(self.bn1_2(self.conv1_2(h))) # h = self.pool1(h) # h = self.relu2_1(self.bn2_1(self.conv2_1(h))) # h = self.relu2_2(self.bn2_2(self.conv2_2(h))) # h = self.pool2(h) # h = self.relu3_1(self.bn3_1(self.conv3_1(h))) # h = self.relu3_2(self.bn3_2(self.conv3_2(h))) # h = self.relu3_3(self.bn3_3(self.conv3_3(h))) # h = self.pool3(h) # pool3 = h # 1/8 # h = self.relu4_1(self.bn4_1(self.conv4_1(h))) # h = self.relu4_2(self.bn4_2(self.conv4_2(h))) # h = self.relu4_3(self.bn4_3(self.conv4_3(h))) # h = self.pool4(h) # pool4 = h # 1/16 # h = self.relu5_1(self.bn5_1(self.conv5_1(h))) # h = self.relu5_2(self.bn5_2(self.conv5_2(h))) # h = self.relu5_3(self.bn5_3(self.conv5_3(h))) # h = self.pool5(h) # h = self.relu6(self.bn6(self.fc6(h))) # h = self.drop6(h) # h = self.relu7(self.bn7(self.fc7(h))) # h = self.drop7(h) # h = self.score_fr(h) # h = self.upscore2(h) # upscore2 = h # 1/16 # h = self.score_pool4(pool4) # # crop for alignment # h = h[:, :, 5:5 + upscore2.size()[2], 5:5 + upscore2.size()[3]] # score_pool4c = h # 1/16 # h = upscore2 + score_pool4c # 1/16 # h = self.upscore_pool4(h) # upscore_pool4 = h # 1/8 # h = self.score_pool3(pool3) # # crop # h = h[:, :, # 9:9 + upscore_pool4.size()[2], # 9:9 + upscore_pool4.size()[3]] # score_pool3c = h # 1/8 # h = upscore_pool4 + score_pool3c # 1/8 # h = self.upscore8(h) # # crop # h = h[:, :, 31:31 + x.size()[2], 31:31 + x.size()[3]].contiguous() # return h