FCN的概念是在15年CVPR的Fully Convolutional Networks for Semantic Segmentation中提出，这篇论文第一次将端到端的卷积网络推广到了语义分割的任务当中，本篇博客主要会介绍该论文的Pytorch版本的简单实现。

介绍

FCN的主要特点：

• 做像素级别的分类，使用ground-truth作为监督信息
• 将全连接层改为卷积层
• 使用反卷积deconv做上采样upsampling,从而得到跟原图大小一致的特征图

• 采用Skip-connection,结合不同深度层的特征图，从而衍生出FCN-32s, FCN-16s, FCN-8s.过程如图1所示。

  
  
  图1 Skip Connection

Pytorch实现

实现部分大概可分为数据集制作、网络结构实现和训练三部分。

数据集制作

本次实验选用了一个比较小的数据集，是一些随机背景上的一些包的图片，数据集可以点击这里获取。

首先我们先定义一个函数来获取每张图片的地址并划分为训练集,测试集

def read_image(val_ratio=0.3, train=True):
   data_dir = os.path.join(cur_path, "bag_data")
   file_name = os.listdir(data_dir)
   img_path = [os.path.join(cur_path, "bag_data", file_name[i]) for i in range(len(file_name))]
   mask_path = [os.path.join(cur_path, "bag_data_msk", file_name[i]) for i in range(len(file_name))]
   # val_ratio设为0.3，代表70%数据用于训练，30%用于验证(这里并没有加上测试集）
   count = int(len(file_name)*val_ratio)
   if train:
       return img_path[count:], mask_path[count:]
   else:
       return img_path[:count], mask_path[:count]

接下来我们则定义自己的dataset类用于后续的训练。

class Bag_dataset(Dataset):
    def __init__(self, train=True, transforms=None, val_ratio=0.3):
        self.transforms = transforms
        self.data_path, self.mask_path = read_image(val_ratio, train)
        print("read " + str(len(self.data_path)) + " images")
    
    def __getitem__(self, item):
        img_path = self.data_path[item]
        mask_path = self.mask_path[item]
        img = Image.open(img_path).convert('RGB')
        mask = Image.open(mask_path).convert('L')
        img, mask = self.transforms(img, mask)
        return img, mask

    def __len__(self):
        return len(self.data_path)

再定义我们的transfroms函数 用于对读入的图片数据进行处理，这里有几个点要注意：

• 本身来说全卷积网络去掉了全连接层，可以支持任意大小的图片输入，但是由于训练时同一个batch的数据大小必须一致，所以我这里对图片统一做了resize操作，不过这样会影响到网络精度。如果为了提高精度所以可以不做resize，只需将batch_size设为1即可.
• ToTensor() 函数是将PIL或者ndarray格式的数据(HxWxC)转为FloatTensor格式的 CxHxW.
• 对mask的处理，mask中目标区域为0，背景区域为255，我们首先将背景区域全部转为-1，再整体加一，这样便得到了0,1标签的二维矩阵，最后将其转为Tensor格式。

def train_transform(img, mask):
    #定义数据处理模块
    im_aug = tfs.Compose([
        
        tfs.Resize((320, 320)),
        # toTensor
        tfs.ToTensor(),
        tfs.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    img = im_aug(img)
    mask = tfs.Resize((320, 320))(mask)
    mask = np.array(mask, dtype=np.int64)
    mask[mask > 0] = -1
    mask += 1
    mask = torch.from_numpy(mask)
    return img, mask

数据集模块代码基本完成，最后做一个接口给外部调用, arg是我程序中的配置类。

def get_dataset(arg, train=True):

   dataset = Bag_dataset(train, train_transform, arg.val_ratio)
   data_loader = DataLoader(dataset, arg.train_batch_size if train else arg.test_batch_size, shuffle=True,
                            num_workers=arg.num_worker)

   return data_loader

网络结构实现

本篇论文的网络结构主要由卷积和反卷积模块构成，前面使用卷积模块进行特征提取，这部分可以使用一些现成的网络结构如VGG，ResNet等，再使用反卷积层将图像恢复到原始大小，最后再使用一个1x1的卷积核进行像素级别的分类。我这里只介绍FCN-8s的实现。

class FCN8s(nn.Module):
   def __init__(self, pretrained_net, num_class=2, verbose=False):
       super().__init__()
       self.num_class = num_class
       self.verbose = verbose
       # 我使用的预训练网络为ResNet-34
       self.stage1 = nn.Sequential(*list(pretrained_net.children()))[:-4]
       self.stage2 = nn.Sequential(*list(pretrained_net.children()))[-4]
       self.stage3 = nn.Sequential(*list(pretrained_net.children()))[-3]

       self.relu = nn.ReLU(True)
       # 定义反卷积层，反卷积层的具体计算可查看pytorch官方手册
       self.deconv1 = nn.ConvTranspose2d(512, 256, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
       self.bn1 = nn.BatchNorm2d(256)
       self.deconv2 = nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
       self.bn2 = nn.BatchNorm2d(128)
       self.deconv3 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
       self.bn3 = nn.BatchNorm2d(64)
       self.deconv4 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
       self.bn4 = nn.BatchNorm2d(32)
       self.deconv5 = nn.ConvTranspose2d(32, 16, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
       self.bn5 = nn.BatchNorm2d(16)
       # 定义1x1的卷积核用于像素分类
       self.classifier = nn.Conv2d(16, num_class, kernel_size=1)

   def forward(self, x):
       
       x3 = self.stage1(x)
       x4 = self.stage2(x3)
       x5 = self.stage3(x4)
       # 使用反卷积做图像像素恢复 将x5的size恢复至于x4一致，后续即可融合x5和x4(Skip Connection)
       score = self.relu(self.deconv1(x5))
       # 使用上采样的目的是当输入图片大小不一致时，可能使用反卷积恢复的x5大小会跟x4相差1个像素，所以采样上采样来补偿(如果有别的办法也可以推荐给我)
       score = F.interpolate(score, x4.size()[2:])
       score = self.bn1(score+x4)
       score = self.relu(self.deconv2(score))
       score = F.interpolate(score, x3.size()[2:])
       score = self.bn2(score+x3)
       score = self.bn3(self.relu(self.deconv3(score)))
       score = self.bn4(self.relu(self.deconv4(score)))
       score = self.bn5(self.relu(self.deconv5(score)))
       score = self.classifier(score)

       return score

训练

首先我们先定义损失函数。因为是分类任务所以我们采用了交叉熵。不过因为输入的mask标签为二维数据，所以我们预先做了一些处理。

def cross_entropy2d(input, target, weight=None, size_average=True):
    # input (n, c, h, w) , target(n, h, w)
    n, c, h, w = input.size()
    log_p = F.log_softmax(input, dim=1)
    log_p = log_p.transpose(1, 2).transpose(2, 3).contiguous()
    log_p = log_p[target.view(n, h, w, 1).repeat(1, 1, 1, c) >= 0]
    log_p = log_p.view(-1, c)
    mask = target >= 0
    target = target[mask]
    loss = F.nll_loss(log_p, target, weight=weight, reduction='sum')
    if size_average:
        loss /= mask.data.sum()
    return loss

为了衡量网络训练的结构，我们再定义一个测量的准确度函数

def get_accuracy(output, target):
   v, i = torch.max(output, 1)
   return torch.sum(target.long() == i).float()/target.numel()

有了这两个函数我们就可以开始进行网络的训练了，具体的训练代码这里就不贴出来了，可以去看我的GitHub源码