FCN的代码解读

目录

模型初始化

VGG初始化

FCN初始化 

图片的预处理

图片处理

图片编码

计算相关参数

模型训练

 一个小问题

完整代码

参考


最近浅研究了一下关于图像领域的图像分割的相关知识,发现水还是挺深的,因为FCN差不多也是领域的开山鼻祖,所以就先从这个方面入手。理论就不多讲很多了,网上一搜一大堆,主要就是解析一下代码部分。

模型初始化

众所周知,FCN的后半段是新的,前半段一般移植自其他模型,这里我选择了Vgg16的模型结构。所以模型初始化分为两步,首先是对Vgg网络的初始化,然后是对Fcn网络的初始化。

VGG初始化

 这里为了方便Vgg的选择将几个不同的Vgg封装在了列表中,以数字代表卷积后的输出通道数,卷积的输入通道数也就是前一个的输出通道数,M代表池化层,Vgg采用的都是卷积核为3的卷积层和大小为2的池化层核,所以这两个参数为已知无需标注。

还需要注意的是因为Fcn是全卷积网络,所以是不需要最后的全连接层的,所以去掉。

# Vgg网络结构配置(数字代表经过卷积后的channel数,‘M’代表池化层)
cfg = {
    'vgg11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'vgg13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'vgg16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'vgg19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}

# 由cfg构建vgg-Net的卷积层和池化层(block1-block5)
def make_layers(cfg, batch_norm=False):
    layers = []
    in_channels = 3  # RGB初始值
    for v in cfg:
        if v == 'M':  # 池化层
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:  # 是否需要归一化
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v  # 这一层输出的通道数就是下一层输入的通道数
    return nn.Sequential(*layers)

# 下面开始构建VGGnet
class VGGNet(VGG):
    def __init__(self, pretrained=True, model='vgg16', requires_grad=True, remove_fc=True, show_params=False):
        super().__init__(make_layers(cfg[model]))
        self.ranges = ranges[model]  # ranges是一个字典,键是model名字,后面的是池化层的信息

        # 获取VGG模型训练好的参数,并加载(第一次执行需要下载一段时间)
        if pretrained:
            exec("self.load_state_dict(models.%s(pretrained=True).state_dict())" % model)

        # 屏蔽预训练模型的权重,只训练最后一层的全连接的权重,因为fcn模型是建立在vgg16基础上训练的,所以前面训练好的VGG网络不修改
        if not requires_grad:
            for param in super().parameters():
                param.requires_grad = False

        # 去掉vgg最后的全连接层(classifier)
        if remove_fc:
            del self.classifier

        # 打印网络的结构
        if show_params == True:
            for name, param in self.named_parameters():
                print(name, param.size())

    def forward(self, x):
        output = {}
        # 利用之前定义的ranges获取每个max-pooling层输出的特征图,这个主要是FCN32的上采样要用到
        for idx, (begin, end) in enumerate(self.ranges):  # enumerate用于枚举,同时给出元素和下标
            # self.ranges = ((0, 5), (5, 10), (10, 17), (17, 24), (24, 31)) (vgg16 examples)
            for layer in range(begin, end):
                x = self.features[layer](x)
                # 相当于把x矩阵放进layer层,然后得到输出,0-5代表第一个max-pool需要经过的层数,所以x1实际上就是第一个max-pool层输出
            output["x%d" % (idx + 1)] = x
            # x数字越大越深
        # output 为一个字典键x1d对应第一个max-pooling输出的特征图,x2...x5类推
        return output

Fcn8s是需要融合前面3个池化层信息的,所以需要将Vgg模型的池化层信息记录下来,这也是foward在做的事情,可以看到这串代码实际上就是取出一串卷积层加上最后的池化层,做完之后把结果存储到字典中,最后output中存储的就是几个池化层的信息(因为每次都是以池化层为结束)。 

FCN初始化 

然后是关于FCN网络的初始化。FCN下有FCN32s,FCN16s,FCN8s,如下图:

FCN的代码解读_第1张图片

这是FCN8s,因为融合了不同深度的池化层的信息,因而相比直接输出对边缘处理会更加丝滑,因为浅的抽象层次往往对细节有着更好理解。但是作者也说了,并不是融合的越多越多好,Fcn4s相比并没有很大的精度提高,因此也是适可而止,因此下面就直接做Fcn8s。

# 下面由VGG构建FCN8s
class FCN8s(nn.Module):

    def __init__(self, pretrained_net, n_class):
        super().__init__()
        # 定义可能会用到的东西
        self.n_class = n_class
        self.pretrained_net = pretrained_net
        self.conv6 = nn.Conv2d(512, 512, kernel_size=1, stride=1, padding=0, dilation=1)
        self.conv7 = nn.Conv2d(512, 512, kernel_size=1, stride=1, padding=0, dilation=1)  # 卷积核大小是1,本质上是全连接层
        # 这里写两个一样的可能是为了写出前后关系的感觉?
        self.relu = nn.ReLU(inplace=True)
        self.deconv1 = nn.ConvTranspose2d(512, 512, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn1 = nn.BatchNorm2d(512)
        self.deconv2 = nn.ConvTranspose2d(512, 256, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn2 = nn.BatchNorm2d(256)
        self.deconv3 = nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.deconv4 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn4 = nn.BatchNorm2d(64)
        self.deconv5 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn5 = nn.BatchNorm2d(32)
        self.classifier = nn.Conv2d(32, n_class, kernel_size=1)

    def forward(self, x):
        output = self.pretrained_net(x)
        # 这个已经在前面的forward中初始化了,里面已经存储了相关特征图
        x5 = output['x5']  # max-pooling5的feature map (1/32)  5*5,160/32
        # print(x5.size())
        x4 = output['x4']  # max-pooling4的feature map (1/16)
        x3 = output['x3']  # max-pooling3的feature map (1/8)

        # 所以总结一下FCN里面的几个合成的步骤也就是反卷积->激活->标准化->加上前面的pool层继续

        score = self.relu(self.conv6(x5))  # conv6  size不变 (1/32)
        # score = self.relu(self.conv7(score))  # conv7  size不变 (1/32)

        # 这里我尝试把右边括号里的x5改成了score
        score = self.relu(self.deconv1(score))  # out_size = 2*in_size (1/16)
        # print(score.size())  # 反卷积之后变为两倍
        score = self.bn1(score + x4)  # bn是标准化,表示加x4第二池化层的结果一同进行计算
        score = self.relu(self.deconv2(score))  # out_size = 2*in_size (1/8)

        score = self.bn2(score + x3)

        # 到这里为止就是全部的FCN步骤,接下来是反卷积到原尺寸

        # 此时是1/8,然后继续反卷积,每次扩大两倍边长直到最后和原图一样
        score = self.bn3(self.relu(self.deconv3(score)))  # out_size = 2*in_size (1/4),反卷积后标准化
        score = self.bn4(self.relu(self.deconv4(score)))  # out_size = 2*in_size (1/2)
        score = self.bn5(self.relu(self.deconv5(score)))  # out_size = 2*in_size (1)
        score = self.classifier(score)  # size不变,使输出的channel等于类别数,相当于对每个点分类
        return score

因为代码是取自其他博主,因而在阅读过程中也遇到了一些问题,原代码对于score的处理如下,但是可以看到第一句和第二句对score处理了之后在第三句又对score重新赋值,这就代表了什么,前两句是无效的,这也是我疑惑的地方,后来我也去参考了一下这位博主参考的githug源码,猜想应该是要把处理后的score放进去继续处理,也就成为了上面的样子。

score = self.relu(self.conv6(x5))    # conv6  size不变 (1/32)
score = self.relu(self.conv7(score)) # conv7  size不变 (1/32)
score = self.relu(self.deconv1(x5))   # out_size = 2*in_size (1/16)       
score = self.bn1(score + x4)                      
score = self.relu(self.deconv2(score)) # out_size = 2*in_size (1/8)           
score = self.bn2(score + x3)                      
score = self.bn3(self.relu(self.deconv3(score)))  # out_size = 2*in_size (1/4)
score = self.bn4(self.relu(self.deconv4(score)))  # out_size = 2*in_size (1/2)
score = self.bn5(self.relu(self.deconv5(score)))  # out_size = 2*in_size (1)
score = self.classifier(score)                    # size不变,使输出的channel等于类别数

每一次池化尺寸会减半,而后面每次反卷积就意味着尺寸会变为两倍,因此处理到最后也就成为了原来的尺寸。

图片的预处理

接下来就是对于训练图片的预处理,包括图片处理和图片编码部分。

图片处理

对图片本身的处理主要是尺寸变换还有标准化和打包这些,基本是通过库函数来完成,就不多说。 

图片编码

编码相对麻烦,需要用到独热编码,因为损失函数计算可能会用到。

独热编码就是开辟n个位置,在对应的那个维度为1,剩下为0。比如性别可以是男/女,男是第一个,女是第二个,那么对于一个个体他的性别可以是男,编码10,或者女,编码01;再假设国籍可以是中国/美国/日本,那么一个人的国籍编码可以是100,010,001(中国,美国,日本),也就是永远一个为1,其他为0,为1的就对应他自己所属的。这里的类别也类似,假设有两个像素点,每个像素点要么01,属于第一类,要么10,属于第二类。

独热编码如下:

def onehot(data, n):
    buf = np.zeros(data.shape + (n,))  # 相当于给每一个像素开辟一个维度,除了他其他都是其他
    nmsk = np.arange(data.size) * n + data.ravel()  # revel表示展平多维数组,就是flatten
    # 前面的data.size是从第一个元素到最后一个元素(所有),下标0--n-1,表示的是行,乘一行个数n就是在在一维数组中一行的开始位置
    buf.ravel()[nmsk] = 1  # 这个就是表示把对应的是1的(根据上面nmsk找到的索引值)值给buf
    return buf

解释一下这个函数是干什么的,传入的参数就是一张图片,比如是160*160的一张图,本质当然是一个数字矩阵,现在要为每个像素点编码,因为有两个类别,所以每个像素点需要两个位置,因此加上一个维度n,全部置0,这个就是没编码前的矩阵,大小160*160*2。nmsk存储的是每一个像素点所对应类别在展平的未编码矩阵中的位置。

举个栗子,现在的图片是二维矩阵([[0,1,0],[1,1,0],[0,0,1]]),那么开辟buf是3*3*2的矩阵,全是0,第一个像素点是0,也就是类别为第一个,所以这个像素点编码[1,0],第二个是1,编码[0,1],后面同理,最后只要把一开始的矩阵中的每个元素换成编码后的就可以了,最终就是[[[1,0],[0,1]......]],但是这样不好写,因此我们可以先把为1的位置记录下来,最后直接替换。展平的编码后的矩阵前四个为1001,我们来讨论怎么来的,第一个像素编码10,而这个1所在最终展平的矩阵中的位置就是0=0*2+0,第二个1所在位置是3=1*2+1,所以可以发现算法:

WZ(1的最终位置)=WZ(像素点索引)*类别数+像素点所属类别

因此就用nmsk将这些1的位置记录下来,然后最后把对应位置的0替换为1,这样就完成了对图像像素的编码。

这样编码后的图片怎么恢复为原来的图,很简单,只要找到1所在的位置是不是就可以了,那是不是就是找最大值在这个维度的位置,也就是argmax()函数,下面是一个简单演示:

    imgB = np.array([1, 0, 1, 1, 0, 1, 1, 0, 0]).reshape(3, 3)
    print('编码前:\n', imgB)
    imgB = onehot(imgB, 2)
    # print('2:', imgB)
    print('恢复:\n', np.argmax(imgB, 2))

效果如下

FCN的代码解读_第2张图片

 这在下面的训练代码中有所体现。

计算相关参数

这里的相关参数指的是精度acc还有iou这些,其他我还没有仔细推算过,主要讲一下精度这个吧。

代码如下:

# 在训练网络前定义函数用于计算Acc 和 mIou
# 计算混淆矩阵
def _fast_hist(label_true, label_pred, n_class):
    mask = (label_true >= 0) & (label_true < n_class)  # 查找有效类别,mask是个bool类型向量
    # 计算匹配个数
    hist = np.bincount(  # bincount输出每个元素的数量,np.bincount([1,1,2]) 输 出 : [0,2,1]代表0有0个,1有2个,2有1个
        n_class * label_true[mask].astype(int) +  # astype代表把bool转为int
        label_pred[mask], minlength=n_class ** 2).reshape(n_class, n_class)  # minlength=4表示最少计算到class*2,为0也计算,不然个数都不够
    '''
    混淆矩阵  n_class = 2,矩阵2*2
        0        1     标答
    0   0*2+0    0*2+1  
    1   1*2+0    1*2+1
    预测
    一维向量的输出是  0,1,2,3,对应到矩阵中
    '''
    return hist


# 根据混淆矩阵计算Acc和mIou
def label_accuracy_score(label_trues, label_preds, n_class):
    """
        Returns accuracy score evaluation result.
      - overall accuracy
      - mean accuracy
      - mean IU
    """
    hist = np.zeros((n_class, n_class))
    for lt, lp in zip(label_trues, label_preds):  # zip(a,b)就是一一对应打包起来
        hist += _fast_hist(lt.flatten(), lp.flatten(), n_class)  # 展平送进去计算,也就是向量计算
    acc = np.diag(hist).sum() / hist.sum()  # 计算主对角线的,也就是正确的数量
    with np.errstate(divide='ignore', invalid='ignore'):
        acc_cls = np.diag(hist) / hist.sum(axis=1)
    acc_cls = np.nanmean(acc_cls)
    with np.errstate(divide='ignore', invalid='ignore'):
        iu = np.diag(hist) / (
                hist.sum(axis=1) + hist.sum(axis=0) - np.diag(hist)
        )
    mean_iu = np.nanmean(iu)
    freq = hist.sum(axis=1) / hist.sum()
    return acc, acc_cls, mean_iu

这里涉及到一个计算混淆矩阵的问题,混淆矩阵本身非常简单,也就是计算00,01,10,11匹配的个数,01代表标注是0,但是预测为1,其他同理。

传入参数就是标答矩阵,预测矩阵和类别数,从注释中可以看出,展开的四个位置0,1,2,3分别是0*2+0,0*2+1,1*2+0,1*2+1,所以这时候将预测值看做行标,标答作为列标就可以很轻松算出0-1匹配情况在四个位置的数量。一开始的musk我猜想是为了剔除无效的坐标,比如预测为3,但实际上没有这个类别也就没有计算的必要了。

至于acc的计算,一定是预测和标答一致才算正确,所以就是对于主对角线求和除以总的像素点个数。

模型训练

以上就是所有相关的轮子,最后开始组装,也就是开始模型训练。

模型训练实际上大同小异,设定优化器,损失函数,然后设定训练轮数,开始训练。

def train(epo_num=50, show_vgg_params=False):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if torch.cuda.is_available():
        print('使用GPU')
    else:
        print('使用CPU')

    vgg_model = VGGNet(requires_grad=True, show_params=show_vgg_params)
    fcn_model = FCN8s(pretrained_net=vgg_model, n_class=2)  # 把训练好的几个maxpool层的集合传给fcn
    fcn_model = fcn_model.to(device)  # 载入模型
    # 这里只有两类,采用二分类常用的损失函数BCE
    criterion = nn.BCELoss().to(device)
    # 随机梯度下降优化,学习率0.001,惯性分数0.7
    optimizer = optim.SGD(fcn_model.parameters(), lr=1e-3, momentum=0.7)

    # 记录训练过程相关指标
    all_train_iter_loss = []
    all_test_iter_loss = []
    test_Acc = []
    test_mIou = []
    # start timing
    prev_time = datetime.now()

    for epo in range(1, epo_num + 1):
        pbar = tqdm(train_dataloader)  # 要先把训练集转进进度条里面
        # 训练
        train_loss = 0  # 一轮的总误差,全部图片的
        fcn_model.train()
        for index, (bag, bag_msk) in enumerate(pbar):
            bag = bag.to(device)
            bag_msk = bag_msk.to(device)

            optimizer.zero_grad()  # 梯度清零
            output = fcn_model(bag)  # 输出
            # print(output.shape)
            output = torch.sigmoid(output)  # output.shape is torch.Size([4, 2, 160, 160])
            loss = criterion(output, bag_msk)  # 计算和标答的误差
            # print('loss=',loss)
            loss.backward()  # 需要计算导数,则调用backward()
            # print('grad_loss=',loss)
            iter_loss = loss.item()  # .item()返回一个具体的值,一般用于loss和acc,这一张的误差
            all_train_iter_loss.append(iter_loss)  # 把误差放进误差列表,方便最后画图
            train_loss += iter_loss  # 加到一轮总的误差里
            optimizer.step()  # 根据求导得到的进行更新

            output_np = output.cpu().detach().numpy().copy()
            bag_msk_np = bag_msk.cpu().detach().numpy().copy()
            bag_msk_np = np.argmax(bag_msk_np, axis=1)

            info = 'epoch {}, {}/{},train loss is {}'.format(epo, index, len(train_dataloader), iter_loss)
            pbar.set_description(info)
        # 验证
        test_loss = 0
        fcn_model.eval()
        with torch.no_grad():
            for index, (bag, bag_msk) in enumerate(test_dataloader):
                bag = bag.to(device)
                bag_msk = bag_msk.to(device)

                optimizer.zero_grad()
                output = fcn_model(bag)
                output = torch.sigmoid(output)  # output.shape is torch.Size([4, 2, 160, 160])
                loss = criterion(output, bag_msk)
                iter_loss = loss.item()
                all_test_iter_loss.append(iter_loss)
                test_loss += iter_loss  # 计算并记录误差

                output_np = output.cpu().detach().numpy().copy()
                output_np = np.argmax(output_np, axis=1)
                bag_msk_np = bag_msk.cpu().detach().numpy().copy()

        # 计算时间
        cur_time = datetime.now()
        # divmod(x,y)返回一个元组,第一个参数是整除的结果,第二个是取模的结果
        h, remainder = divmod((cur_time - prev_time).seconds, 3600)
        m, s = divmod(remainder, 60)
        time_str = "Time %02d:%02d:%02d" % (h, m, s)  # 时分秒
        prev_time = cur_time  # 更新时间
        info = 'epoch: %d, epoch train loss = %f, epoch test loss = %f, %s' \
               % (epo, train_loss / len(train_dataloader), test_loss / len(test_dataloader), time_str)
        print(info)

        acc, acc_cls, mean_iu = label_accuracy_score(bag_msk_np, output_np, 2)
        test_Acc.append(acc)
        test_mIou.append(mean_iu)

        print('Acc = %f, mIou = %f' % (acc, mean_iu))
        # 每2个epoch存储一次模型
        if np.mod(epo, 2) == 0:
            # 只存储模型参数
            torch.save(fcn_model.state_dict(), './pths/fcn_model_{}.pth'.format(epo))
            print('成功存储模型:fcn_model_{}.pth'.format(epo))

 一个小问题

正文在上面就结束了,但是我还是有一个疑问,除了上面FCN模型那里有点小问题,还有一个地方就是关于onehot()中nmsk的计算,原作者的代码如下:

def onehot(data, n):
    buf = np.zeros(data.shape + (n,))  # 相当于给每一个像素开辟一个维度,除了他其他都是其他
    nmsk = np.arange(data.size) * n + data.ravel()  # revel表示展平多维数组,就是flatten
    buf.ravel()[nmsk-1] = 1  # 这个就是表示把对应的是1的(根据上面nmsk找到的索引值)值给buf
    return buf

区别就是这里的nmsk有一个-1,并且在恢复矩阵时选择了argmin()而非argmax()函数,但是实际上我用这样的一套去编码一个3*3矩阵在还原时,矩阵已经变样。

 FCN的代码解读_第3张图片

 可以看到无法恢复,但是奇怪的是我用这样的规则去看了恢复的图(下图中中间是标注,左边是用了nmsk-1和argmin()的组合,右边是nmsk和argmax()的组合)

FCN的代码解读_第4张图片

竟然毫无违和恢复了。嗯???还有这种操作?好像也没什么问题。众所周知,为什么可以比为什么不可以更加离奇。我百思不得其解,后来想了想这也许和图片本身一些特殊的性质也有关系,大致如下。

这种图片首先是二分类,非1即0,所以这也就给了找最大1变为找最小也就是找0,使用argmin()的机会,那么按说这时候的输出应该是黑白颠倒,但是实际上并没有,为什么?因为nmsk-1。

假设编码对象是111000,那么正常编码后展平就应该是01 01 01 10 10 10,但是由于nmsk-1了,所以所有1的位置都要前移,第一个变成-1,到了最后,最后编码结果为10 10 11 01 01 00,然后这时候两个相邻之间的最小值索引发现是1,1,0,0,0,0,可以发现两点:

是大部分正常恢复了,为什么,因为0101..前移之后变成1010...然后找最小,0代替了原来的1,所以现在的找最小等同于原来的找最大。

是第三个1恢复出错了,为什么,这是由于移动导致两个01编码后本来是0110,然后11都移动到了1所对应的位置,然后argmin()对于相同的参数输出第一个索引下标,就成为了0,恢复出错,从上面的程序结果来看也是如此,对于每一个10的交界处,1都被恢复为了0,导致出错。

那么为什么上图恢复出来看上去没有什么问题呢?答案就是一张图片10交界太少了,大部分都是000....111....000....111...,导致这种错误发生的那几个像素点几乎不影响最终结果。

完整代码

import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset, random_split
from torchvision import transforms
from torchvision.models.vgg import VGG
import cv2
import numpy as np
from tqdm import tqdm


# 将标记图(每个像素值代该位置像素点的类别)转换为onehot编码
def onehot(data, n):
    buf = np.zeros(data.shape + (n,))  # 相当于给每一个像素开辟一个维度,除了他其他都是其他
    nmsk = np.arange(data.size) * n + data.ravel()  # revel表示展平多维数组,就是flatten
    # 前面的data.size是从第一个元素到最后一个元素(所有),下标0--n-1,表示的是行,乘一行个数n就是在在一维数组中一行的开始位置
    # 后面的是0--n-1表示的是类别,表示第几个
    # 索引nmsk存储了在一维数组中应该是1的位置,也就是正确答案
    buf.ravel()[nmsk-1] = 1  # 这个就是表示把对应的是1的(根据上面nmsk找到的索引值)值给buf
    return buf


# 利用torchvision提供的transform,定义原始图片的预处理步骤(转换为tensor和标准化处理)
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

# 利用torch提供的Dataset类,定义我们自己的数据集
base_img = './data/bag_data/'  # 训练集地址
base_img_msk = './data/bag_data_msk/'  # 标注地址


class BagDataset(Dataset):

    def __init__(self, transform=None):
        self.transform = transform

    def __len__(self):
        return len(os.listdir(base_img))

    def __getitem__(self, idx):
        img_name = os.listdir(base_img)[idx]  # index是随机数,是图片的索引值
        imgA = cv2.imread(base_img + img_name)
        imgA = cv2.resize(imgA, (160, 160))
        # img_name = '1.jpg'
        imgB = cv2.imread(base_img_msk + img_name, 0)
        imgB = cv2.resize(imgB, (160, 160))
        # 下面是对标注的一些处理
        imgB = imgB / 255  # 归一化
        imgB = imgB.astype('uint8')  # 转化成整数
        imgB = onehot(imgB, 2)
        imgB = imgB.transpose(2, 0, 1)  # 转置  0 1 2 -> 2 0 1 相当于几个维度的位置关系变化,就是把一开始加到最后的提到最前面,效果就是把两列的每一列变成一张图
        imgB = torch.FloatTensor(imgB)

        if self.transform:
            imgA = self.transform(imgA)

        return imgA, imgB


# 实例化数据集
bag = BagDataset(transform)

train_size = int(0.9 * len(bag))
test_size = len(bag) - train_size
train_dataset, test_dataset = random_split(bag, [train_size, test_size])  # 划分数据集

# 利用DataLoader生成一个分batch获取数据的可迭代对象
train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=4)
test_dataloader = DataLoader(test_dataset, batch_size=4, shuffle=True, num_workers=4)

# <-------------------------------------------------------->#
# 下面开始定义网络模型
# 先定义VGG结构

# ranges 是用于方便获取和记录每个池化层得到的特征图
# 例如vgg16,需要(0, 5)的原因是为方便记录第一个pooling层得到的输出(详见下午、稳VGG定义)
ranges = {
    'vgg11': ((0, 3), (3, 6), (6, 11), (11, 16), (16, 21)),
    'vgg13': ((0, 5), (5, 10), (10, 15), (15, 20), (20, 25)),
    'vgg16': ((0, 5), (5, 10), (10, 17), (17, 24), (24, 31)),
    'vgg19': ((0, 5), (5, 10), (10, 19), (19, 28), (28, 37))
}

# Vgg网络结构配置(数字代表经过卷积后的channel数,‘M’代表池化层)
cfg = {
    'vgg11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'vgg13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'vgg16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'vgg19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}


# 由cfg构建vgg-Net的卷积层和池化层(block1-block5)
def make_layers(cfg, batch_norm=False):
    layers = []
    in_channels = 3  # RGB初始值
    for v in cfg:
        if v == 'M':  # 池化层
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:  # 是否需要归一化
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v  # 这一层输出的通道数就是下一层输入的通道数
    return nn.Sequential(*layers)


# 下面开始构建VGGnet
class VGGNet(VGG):
    def __init__(self, pretrained=True, model='vgg16', requires_grad=True, remove_fc=True, show_params=False):
        super().__init__(make_layers(cfg[model]))
        self.ranges = ranges[model]  # ranges是一个字典,键是model名字,后面的是池化层的信息

        # 获取VGG模型训练好的参数,并加载(第一次执行需要下载一段时间)
        if pretrained:
            exec("self.load_state_dict(models.%s(pretrained=True).state_dict())" % model)

        # 屏蔽预训练模型的权重,只训练最后一层的全连接的权重,因为fcn模型是建立在vgg16基础上训练的,所以前面训练好的VGG网络不修改
        if not requires_grad:
            for param in super().parameters():
                param.requires_grad = False

        # 去掉vgg最后的全连接层(classifier)
        if remove_fc:
            del self.classifier

        # 打印网络的结构
        if show_params == True:
            for name, param in self.named_parameters():
                print(name, param.size())

    def forward(self, x):
        output = {}
        # 利用之前定义的ranges获取每个max-pooling层输出的特征图,这个主要是FCN32的上采样要用到
        for idx, (begin, end) in enumerate(self.ranges):  # enumerate用于枚举,同时给出元素和下标
            # self.ranges = ((0, 5), (5, 10), (10, 17), (17, 24), (24, 31)) (vgg16 examples)
            for layer in range(begin, end):
                x = self.features[layer](x)
                # 相当于把x矩阵放进layer层,然后得到输出,0-5代表第一个max-pool需要经过的层数,所以x1实际上就是第一个max-pool层输出
            output["x%d" % (idx + 1)] = x
            # x数字越大越深
        # output 为一个字典键x1d对应第一个max-pooling输出的特征图,x2...x5类推
        return output



# 下面由VGG构建FCN8s
class FCN8s(nn.Module):

    def __init__(self, pretrained_net, n_class):
        super().__init__()
        # 定义可能会用到的东西
        self.n_class = n_class
        self.pretrained_net = pretrained_net
        self.conv6 = nn.Conv2d(512, 512, kernel_size=1, stride=1, padding=0, dilation=1)
        self.conv7 = nn.Conv2d(512, 512, kernel_size=1, stride=1, padding=0, dilation=1)  # 卷积核大小是1,本质上是全连接层
        # 这里写两个一样的可能是为了写出前后关系的感觉?
        self.relu = nn.ReLU(inplace=True)
        self.deconv1 = nn.ConvTranspose2d(512, 512, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn1 = nn.BatchNorm2d(512)
        self.deconv2 = nn.ConvTranspose2d(512, 256, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn2 = nn.BatchNorm2d(256)
        self.deconv3 = nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.deconv4 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn4 = nn.BatchNorm2d(64)
        self.deconv5 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn5 = nn.BatchNorm2d(32)
        self.classifier = nn.Conv2d(32, n_class, kernel_size=1)

    def forward(self, x):
        output = self.pretrained_net(x)
        # 这个已经在前面的forward中初始化了,里面已经存储了相关特征图
        x5 = output['x5']  # max-pooling5的feature map (1/32)  5*5,160/32
        x4 = output['x4']  # max-pooling4的feature map (1/16)
        x3 = output['x3']  # max-pooling3的feature map (1/8)

        # 所以总结一下FCN里面的几个合成的步骤也就是反卷积->激活->标准化->加上前面的pool层继续

        # 这两句没用,或者说用错了
        score = self.relu(self.conv6(x5))  # conv6  size不变 (1/32)

        # 1/32可能没有融合进去?

        # 这里我尝试把右边括号里的x5改成了score
        score = self.relu(self.deconv1(score))  # out_size = 2*in_size (1/16)

        score = self.bn1(score + x4)  # bn是标准化,表示加x4第二池化层的结果一同进行计算
        score = self.relu(self.deconv2(score))  # out_size = 2*in_size (1/8)

        score = self.bn2(score + x3)

        # 到这里为止就是全部的FCN步骤,接下来是反卷积到原尺寸

        # 此时是1/8,然后继续反卷积,每次扩大两倍边长直到最后和原图一样
        score = self.bn3(self.relu(self.deconv3(score)))  # out_size = 2*in_size (1/4),反卷积后标准化
        score = self.bn4(self.relu(self.deconv4(score)))  # out_size = 2*in_size (1/2)
        score = self.bn5(self.relu(self.deconv5(score)))  # out_size = 2*in_size (1)
        score = self.classifier(score)  # size不变,使输出的channel等于类别数,相当于对每个点分类
        # print(score.shape)
        # time.sleep(1000)
        return score


# <---------------------------------------------->
# 下面开始训练网络

# 在训练网络前定义函数用于计算Acc 和 mIou
# 计算混淆矩阵
def _fast_hist(label_true, label_pred, n_class):
    mask = (label_true >= 0) & (label_true < n_class)  # 查找有效类别,mask是个bool类型向量
    # 计算匹配个数
    hist = np.bincount(  # bincount输出每个元素的数量,np.bincount([1,1,2]) 输 出 : [0,2,1]代表0有0个,1有2个,2有1个
        n_class * label_true[mask].astype(int) +  # astype代表把bool转为int
        label_pred[mask], minlength=n_class ** 2).reshape(n_class, n_class)  # minlength=4表示最少计算到class*2,为0也计算,不然个数都不够
    '''
    混淆矩阵  n_class = 2,矩阵2*2
        0        1     标答
    0   0*2+0    0*2+1  
    1   1*2+0    1*2+1
    预测
    一维向量的输出是  0,1,2,3,对应到矩阵中
    '''
    return hist


# 根据混淆矩阵计算Acc和mIou
def label_accuracy_score(label_trues, label_preds, n_class):
    """
        Returns accuracy score evaluation result.
      - overall accuracy
      - mean accuracy
      - mean IU
    """
    hist = np.zeros((n_class, n_class))
    for lt, lp in zip(label_trues, label_preds):  # zip(a,b)就是一一对应打包起来
        hist += _fast_hist(lt.flatten(), lp.flatten(), n_class)  # 展平送进去计算,也就是向量计算
    acc = np.diag(hist).sum() / hist.sum()  # 计算主对角线的,也就是正确的数量
    with np.errstate(divide='ignore', invalid='ignore'):
        acc_cls = np.diag(hist) / hist.sum(axis=1)
    acc_cls = np.nanmean(acc_cls)
    with np.errstate(divide='ignore', invalid='ignore'):
        iu = np.diag(hist) / (
                hist.sum(axis=1) + hist.sum(axis=0) - np.diag(hist)
        )
    mean_iu = np.nanmean(iu)
    freq = hist.sum(axis=1) / hist.sum()
    return acc, acc_cls, mean_iu


from datetime import datetime

import torch.optim as optim
import matplotlib.pyplot as plt


def train(epo_num=50, show_vgg_params=False):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if torch.cuda.is_available():
        print('使用GPU')
    else:
        print('使用CPU')

    vgg_model = VGGNet(requires_grad=True, show_params=show_vgg_params)
    fcn_model = FCN8s(pretrained_net=vgg_model, n_class=2)  # 把训练好的几个maxpool层的集合传给fcn
    fcn_model = fcn_model.to(device)  # 载入模型
    # 这里只有两类,采用二分类常用的损失函数BCE
    criterion = nn.BCELoss().to(device)
    # 随机梯度下降优化,学习率0.001,惯性分数0.7
    optimizer = optim.SGD(fcn_model.parameters(), lr=1e-3, momentum=0.7)

    # 记录训练过程相关指标
    all_train_iter_loss = []
    all_test_iter_loss = []
    test_Acc = []
    test_mIou = []
    # start timing
    prev_time = datetime.now()

    for epo in range(1, epo_num + 1):
        pbar = tqdm(train_dataloader)  # 要先把训练集转进进度条里面
        # 训练
        train_loss = 0  # 一轮的总误差,全部图片的
        fcn_model.train()
        for index, (bag, bag_msk) in enumerate(pbar):

            bag = bag.to(device)
            bag_msk = bag_msk.to(device)

            optimizer.zero_grad()  # 梯度清零
            output = fcn_model(bag)  # 输出
            # print(output.shape)
            # time.sleep(1000)
            output = torch.sigmoid(output)  # output.shape is torch.Size([4, 2, 160, 160])
            loss = criterion(output, bag_msk)  # 计算和标答的误差
            # print('loss=',loss)
            loss.backward()  # 需要计算导数,则调用backward()
            # print('grad_loss=',loss)
            iter_loss = loss.item()  # .item()返回一个具体的值,一般用于loss和acc,这一张的误差
            all_train_iter_loss.append(iter_loss)  # 把误差放进误差列表,方便最后画图
            train_loss += iter_loss  # 加到一轮总的误差里
            optimizer.step()  # 根据求导得到的进行更新

            output_np = output.cpu().detach().numpy().copy()
            output_np = np.argmax(output_np, axis=1)  # 找出所有通道里面的最小值
            # 相当于就是把两个维度的最小值的找到作为输出,也就是找的是0在两个索引中的位置,本质也是在找1的位置
            bag_msk_np = bag_msk.cpu().detach().numpy().copy()
            bag_msk_np = np.argmax(bag_msk_np, axis=1)
            info = 'epoch {}, {}/{},train loss is {}'.format(epo, index, len(train_dataloader), iter_loss)
            pbar.set_description(info)

        # 验证
        test_loss = 0
        fcn_model.eval()
        with torch.no_grad():
            for index, (bag, bag_msk) in enumerate(test_dataloader):
                bag = bag.to(device)
                bag_msk = bag_msk.to(device)

                optimizer.zero_grad()
                output = fcn_model(bag)
                output = torch.sigmoid(output)  # output.shape is torch.Size([4, 2, 160, 160])
                loss = criterion(output, bag_msk)
                iter_loss = loss.item()
                all_test_iter_loss.append(iter_loss)
                test_loss += iter_loss  # 计算并记录误差

                output_np = output.cpu().detach().numpy().copy()
                output_np = np.argmax(output_np, axis=1)
                bag_msk_np = bag_msk.cpu().detach().numpy().copy()
                # 解释一下为什么这里的0和1一样多,因为按照onehot,这里一开始实际上每个像素点对应onehot变化是[0,1]或者[1,0],所以10的总和是一样,因为每个像素点对应了一组[1,0]
                # 之后经过一个维度变换,160,160,2-->2,160,160也就是被分成了两张图片,找两个维度0所在的索引
                bag_msk_np = np.argmax(bag_msk_np, axis=1)

        # 计算时间
        cur_time = datetime.now()
        # divmod(x,y)返回一个元组,第一个参数是整除的结果,第二个是取模的结果
        h, remainder = divmod((cur_time - prev_time).seconds, 3600)
        m, s = divmod(remainder, 60)
        time_str = "Time %02d:%02d:%02d" % (h, m, s)  # 时分秒
        prev_time = cur_time  # 更新时间
        # print()
        info = 'epoch: %d, epoch train loss = %f, epoch test loss = %f, %s' \
               % (epo, train_loss / len(train_dataloader), test_loss / len(test_dataloader), time_str)
        print(info)

        acc, acc_cls, mean_iu = label_accuracy_score(bag_msk_np, output_np, 2)
        test_Acc.append(acc)
        test_mIou.append(mean_iu)

        print('Acc = %f, mIou = %f' % (acc, mean_iu))
        # 每2个epoch存储一次模型
        if np.mod(epo, 2) == 0:
            # 只存储模型参数
            torch.save(fcn_model.state_dict(), './pths/fcn_model_{}.pth'.format(epo))
            print('成功存储模型:fcn_model_{}.pth'.format(epo))
    # 绘制训练过程数据
    plt.figure()
    plt.subplot(221)
    plt.title('train_loss')
    plt.plot(all_train_iter_loss)
    plt.xlabel('batch')
    plt.subplot(222)
    plt.title('test_loss')
    plt.plot(all_test_iter_loss)
    plt.xlabel('batch')
    plt.subplot(223)
    plt.title('test_Acc')
    plt.plot(test_Acc)
    plt.xlabel('epoch')
    plt.subplot(224)
    plt.title('test_mIou')
    plt.plot(test_mIou)
    plt.xlabel('epoch')
    plt.show()


if __name__ == "__main__":
    # 主程序
    train(epo_num=20, show_vgg_params=False)  # 参数是设置是否打印网络结构

参考

代码来源

FCN详解与pytorch简单实现(附详细代码解读)_zinc_abc的博客-CSDN博客

数据集和代码

mirrors / bat67 / pytorch-FCN-easiest-demo · GitCode

你可能感兴趣的:(深度学习,神经网络,人工智能)