Pytorch 分割模型构建和训练【直播】2019 年县域农业大脑AI挑战赛---(四)模型构建和网络训练

对于分割网络，如果当成一个黑箱就是：输入一个3x1024x1024 输出4x1024x1024。

我没有使用二分类，直接使用了四分类。

分类网络使用了SegNet，没有加载预训练模型，参数也是默认初始化。为了加快训练，1024输入进网络后直接通过

pooling缩小到256的尺寸，等到输出层，直接使用bilinear放大4倍，相当于直接在256的尺寸上训练。

import os
import urllib
import torch
import torch.nn as nn
import torch.nn.functional as F

#import torch.utils.model_zoo as model_zoo
from torchvision import models
#https://raw.githubusercontent.com/delta-onera/delta_tb/master/deltatb/networks/net_segnet_bn_relu.py
class SegNet_BN_ReLU(nn.Module):
    # Unet network
    @staticmethod
    def weight_init(m):
        if isinstance(m, nn.Linear):
            torch.nn.init.kaiming_normal(m.weight.data)

    def __init__(self, in_channels, out_channels):
        super(SegNet_BN_ReLU, self).__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels

        self.pool = nn.MaxPool2d(2, return_indices=True)
        self.unpool = nn.MaxUnpool2d(2)

        self.conv1_1 = nn.Conv2d(in_channels, 64, 3, padding=1)
        self.conv1_1_bn = nn.BatchNorm2d(64)
        self.conv1_2 = nn.Conv2d(64, 64, 3, padding=1)
        self.conv1_2_bn = nn.BatchNorm2d(64)

        self.conv2_1 = nn.Conv2d(64, 128, 3, padding=1)
        self.conv2_1_bn = nn.BatchNorm2d(128)
        self.conv2_2 = nn.Conv2d(128, 128, 3, padding=1)
        self.conv2_2_bn = nn.BatchNorm2d(128)

        self.conv3_1 = nn.Conv2d(128, 256, 3, padding=1)
        self.conv3_1_bn = nn.BatchNorm2d(256)
        self.conv3_2 = nn.Conv2d(256, 256, 3, padding=1)
        self.conv3_2_bn = nn.BatchNorm2d(256)
        self.conv3_3 = nn.Conv2d(256, 256, 3, padding=1)
        self.conv3_3_bn = nn.BatchNorm2d(256)

        self.conv4_1 = nn.Conv2d(256, 512, 3, padding=1)
        self.conv4_1_bn = nn.BatchNorm2d(512)
        self.conv4_2 = nn.Conv2d(512, 512, 3, padding=1)
        self.conv4_2_bn = nn.BatchNorm2d(512)
        self.conv4_3 = nn.Conv2d(512, 512, 3, padding=1)
        self.conv4_3_bn = nn.BatchNorm2d(512)

        self.conv5_1 = nn.Conv2d(512, 512, 3, padding=1)
        self.conv5_1_bn = nn.BatchNorm2d(512)
        self.conv5_2 = nn.Conv2d(512, 512, 3, padding=1)
        self.conv5_2_bn = nn.BatchNorm2d(512)
        self.conv5_3 = nn.Conv2d(512, 512, 3, padding=1)
        self.conv5_3_bn = nn.BatchNorm2d(512)

        self.conv5_3_D = nn.Conv2d(512, 512, 3, padding=1)
        self.conv5_3_D_bn = nn.BatchNorm2d(512)
        self.conv5_2_D = nn.Conv2d(512, 512, 3, padding=1)
        self.conv5_2_D_bn = nn.BatchNorm2d(512)
        self.conv5_1_D = nn.Conv2d(512, 512, 3, padding=1)
        self.conv5_1_D_bn = nn.BatchNorm2d(512)

        self.conv4_3_D = nn.Conv2d(512, 512, 3, padding=1)
        self.conv4_3_D_bn = nn.BatchNorm2d(512)
        self.conv4_2_D = nn.Conv2d(512, 512, 3, padding=1)
        self.conv4_2_D_bn = nn.BatchNorm2d(512)
        self.conv4_1_D = nn.Conv2d(512, 256, 3, padding=1)
        self.conv4_1_D_bn = nn.BatchNorm2d(256)

        self.conv3_3_D = nn.Conv2d(256, 256, 3, padding=1)
        self.conv3_3_D_bn = nn.BatchNorm2d(256)
        self.conv3_2_D = nn.Conv2d(256, 256, 3, padding=1)
        self.conv3_2_D_bn = nn.BatchNorm2d(256)
        self.conv3_1_D = nn.Conv2d(256, 128, 3, padding=1)
        self.conv3_1_D_bn = nn.BatchNorm2d(128)

        self.conv2_2_D = nn.Conv2d(128, 128, 3, padding=1)
        self.conv2_2_D_bn = nn.BatchNorm2d(128)
        self.conv2_1_D = nn.Conv2d(128, 64, 3, padding=1)
        self.conv2_1_D_bn = nn.BatchNorm2d(64)

        self.conv1_2_D = nn.Conv2d(64, 64, 3, padding=1)
        self.conv1_2_D_bn = nn.BatchNorm2d(64)
        self.conv1_1_D = nn.Conv2d(64, out_channels, 3, padding=1)

        self.apply(self.weight_init)

    def forward(self, x):
        # Encoder block 1
        x =F.avg_pool2d(x,4)
        #print(x.shape)
        x = self.conv1_1_bn(F.relu(self.conv1_1(x)))
        x1 = self.conv1_2_bn(F.relu(self.conv1_2(x)))
        size1 = x.size()
        x, mask1 = self.pool(x1)

        # Encoder block 2
        x = self.conv2_1_bn(F.relu(self.conv2_1(x)))
        #x = self.drop2_1(x)
        x2 = self.conv2_2_bn(F.relu(self.conv2_2(x)))
        size2 = x.size()
        x, mask2 = self.pool(x2)

        # Encoder block 3
        x = self.conv3_1_bn(F.relu(self.conv3_1(x)))
        x = self.conv3_2_bn(F.relu(self.conv3_2(x)))
        x3 = self.conv3_3_bn(F.relu(self.conv3_3(x)))
        size3 = x.size()
        x, mask3 = self.pool(x3)

        # Encoder block 4
        x = self.conv4_1_bn(F.relu(self.conv4_1(x)))
        x = self.conv4_2_bn(F.relu(self.conv4_2(x)))
        x4 = self.conv4_3_bn(F.relu(self.conv4_3(x)))
        size4 = x.size()
        x, mask4 = self.pool(x4)

        # Encoder block 5
        x = self.conv5_1_bn(F.relu(self.conv5_1(x)))
        x = self.conv5_2_bn(F.relu(self.conv5_2(x)))
        x = self.conv5_3_bn(F.relu(self.conv5_3(x)))
        size5 = x.size()
        x, mask5 = self.pool(x)

        # Decoder block 5
        x = self.unpool(x, mask5, output_size = size5)
        x = self.conv5_3_D_bn(F.relu(self.conv5_3_D(x)))
        x = self.conv5_2_D_bn(F.relu(self.conv5_2_D(x)))
        x = self.conv5_1_D_bn(F.relu(self.conv5_1_D(x)))

        # Decoder block 4
        x = self.unpool(x, mask4, output_size = size4)
        x = self.conv4_3_D_bn(F.relu(self.conv4_3_D(x)))
        x = self.conv4_2_D_bn(F.relu(self.conv4_2_D(x)))
        x = self.conv4_1_D_bn(F.relu(self.conv4_1_D(x)))

        # Decoder block 3
        x = self.unpool(x, mask3, output_size = size3)
        x = self.conv3_3_D_bn(F.relu(self.conv3_3_D(x)))
        x = self.conv3_2_D_bn(F.relu(self.conv3_2_D(x)))
        x = self.conv3_1_D_bn(F.relu(self.conv3_1_D(x)))

        # Decoder block 2
        x = self.unpool(x, mask2, output_size = size2)
        x = self.conv2_2_D_bn(F.relu(self.conv2_2_D(x)))
        x = self.conv2_1_D_bn(F.relu(self.conv2_1_D(x)))

        # Decoder block 1
        x = self.unpool(x, mask1, output_size = size1)
        x = self.conv1_2_D_bn(F.relu(self.conv1_2_D(x)))
        x = self.conv1_1_D(x)
        #print(x.shape)
        return F.interpolate(x,mode='bilinear',scale_factor=4)  

    def load_pretrained_weights(self):

        #vgg16_weights = model_zoo.load_url("https://download.pytorch.org/models/vgg16_bn-6c64b313.pth")
        vgg16_weights=models.vgg16_bn(True).state_dict()
        count_vgg = 0
        count_this = 0

        vggkeys = list(vgg16_weights.keys())
        thiskeys  = list(self.state_dict().keys())

        corresp_map = []

        while(True):
            vggkey = vggkeys[count_vgg]
            thiskey = thiskeys[count_this]

            if "classifier" in vggkey:
                break

            while vggkey.split(".")[-1] not in thiskey:
                count_this += 1
                thiskey = thiskeys[count_this]

            corresp_map.append([vggkey, thiskey])
            count_vgg+=1
            count_this += 1

        mapped_weights = self.state_dict()
        for k_vgg, k_segnet in corresp_map:
            if (self.in_channels != 3) and "features" in k_vgg and "conv1_1." not in k_segnet:
                mapped_weights[k_segnet] = vgg16_weights[k_vgg]
            elif (self.in_channels == 3) and "features" in k_vgg:
                mapped_weights[k_segnet] = vgg16_weights[k_vgg]

        try:
            self.load_state_dict(mapped_weights)
            print("Loaded VGG-16 weights in Segnet !")
        except:
            print("Error VGG-16 weights in Segnet !")
            raise

    def load_from_filename(self, model_path):
        """Load weights from filename."""
        th = torch.load(model_path)  # load the weigths
        self.load_state_dict(th)

def segnet_bn_relu(in_channels, out_channels, pretrained=False, **kwargs):
    """Constructs a ResNet-34 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = SegNet_BN_ReLU(in_channels, out_channels)
    if pretrained:
        model.load_pretrained_weights()
    return model

if __name__=='__main__':
    net=segnet_bn_relu(3,4,False)
    print(net)
    x=torch.rand((1,3,1024,1024))
    print(net.forward(x).shape)

　　训练网络的代码：

import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from farmdataset import FarmDataset

from segnet import segnet_bn_relu as Unet

import time

from PIL import Image

def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        #print(target.shape)
        optimizer.zero_grad()
        output = model(data)
        #print('output size',output.size(),output)

        output = F.log_softmax(output, dim=1)
        loss=nn.NLLLoss2d(weight=torch.Tensor([0.1,0.5,0.5,0.2]).to('cuda'))(output,target)
        loss.backward()

        optimizer.step()

        #time.sleep(0.6)#make gpu sleep
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
    if epoch%2==0:
        imgd=output.detach()[0,:,:,:].cpu()
        img=torch.argmax(imgd,0).byte().numpy()
        imgx=Image.fromarray(img).convert('L')
        imgxx=Image.fromarray(target.detach()[0,:,:].cpu().byte().numpy()*255).convert('L')
        imgx.save("./tmp/predict{}.bmp".format(epoch))
        imgxx.save('./tmp/real{}.bmp'.format(epoch))

def test(args, model, device, testdataset,issave=False):
    model.eval()
    test_loss = 0
    correct = 0
    evalid=[i+7 for i in range(0,2100,15)]
    maxbatch=len(evalid)
    with torch.no_grad():
        for idx in evalid:
            data, target=testdataset[idx]
            data, target = data.unsqueeze(0).to(device), target.unsqueeze(0).to(device)
            #print(target.shape)
            target=target[:,:1472,:1472]
            output = model(data[:,:,:1472,:1472])
            output = F.log_softmax(output, dim=1)
            loss=nn.NLLLoss2d().to('cuda')(output,target)
            test_loss+=loss

            r=torch.argmax(output[0],0).byte()

            tg=target.byte().squeeze(0)
            tmp=0
            count=0
            for i in range(1,4):
                mp=r==i
                tr=tg==i
                tp=mp*tr==1
                t=(mp+tr-tp).sum().item()
                if t==0:
                    continue
                else:
                    tmp+=tp.sum().item()/t
                    count+=1
            if count>0:
                correct+=tmp/count

            if issave:
                Image.fromarray(r.cpu().numpy()).save('predict.png')
                Image.fromarray(tg.cpu().numpy()).save('target.png')
                input()

    print('Test Loss is {:.6f}, mean precision is: {:.4f}%'.format(test_loss/maxbatch,correct))

def main():
    # Training settings
    parser = argparse.ArgumentParser(description='Scratch segmentation Example')
    parser.add_argument('--batch-size', type=int, default=8, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=8, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=30, metavar='N',
                        help='number of epochs to train (default: 10)')
    parser.add_argument('--lr', type=float, default=0.001, metavar='LR',
                        help='learning rate (default: 0.01)')
    parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                        help='SGD momentum (default: 0.5)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='how many batches to wait before logging training status')
    args = parser.parse_args()
    use_cuda = not args.no_cuda and torch.cuda.is_available()

    torch.manual_seed(args.seed)

    device = torch.device("cuda" if use_cuda else "cpu")
    print('my device is :',device)

    kwargs = {'num_workers': 0, 'pin_memory': True} if use_cuda else {}
    train_loader = torch.utils.data.DataLoader(     FarmDataset(istrain=True),batch_size=args.batch_size, shuffle=True,drop_last=True, **kwargs)

    startepoch=0
    model =torch.load('./tmp/model{}'.format(startepoch))  if startepoch else Unet(3,4).to(device)
    args.epochs=50
    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

    for epoch in range(startepoch, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        if epoch %3==0:
            print(epoch)
            test(args, model, device, FarmDataset(istrain=True,isaug=False),issave=False)
            torch.save(model,'./tmp/model{}'.format(epoch))

if __name__ == '__main__':
    main()

　　训练代码每隔三轮，评测一次训练精度，测试数据仍然使用训练数据，只是抽样了。可以根据该精度选择使用什么时刻的网络作为预测节点。

训练精度可以达到0.4，但是这时候貌似过学习了。提交结果并不好。

到现在感觉，只要改进一点，分数就会高一点。如果要继续提高成绩，感觉可以从以下几个方面改进：

样本的不均衡

损失函数

模型结构的设计 可以参考PSP，UNET，deeplab，或者GAN的pix2pix。

总之，感觉只要进行一点改进，功夫就不会白费。

整个从数据切割，数据集准备，数据增强，预测结果保存，深度分割网络 和网络训练，全部代码到此分享完毕，

做完这些你的结果就能到0.2以上。 也是折腾了好几天才到现在，希望这能成为一个基线，看到更精彩的模型思路。

（完）