Pytorch实战2：Carvana Image Masking使用UNet

导包

import os

import numpy as np
import collections
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision

from torch.utils.data import Dataset, DataLoader
from torchvision import  transforms
import torch.optim as optim
import PIL
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split

UNet

• 见Pytorch实战1

设置多卡训练环境

os.environ['CUDA_VISIBLE_DEVICES'] = '4, 5, 6, 7'

• 详情看如何多卡

加载数据

• 和实战1不一样

自行构建Dataset类

# 自行构建Dataset类
class CarvanaDataset(Dataset):
    def __init__(self, base_dir, idx_list, mode='train', transform=None):
        self.base_dir = base_dir
        self.idx_list = idx_list
        # 拼接图片路径
        # listdir：列出文件夹下面的所有文件名（包括文件夹）
        # 对于大量数据这样处理
        # 本次数据集image都是这样命名的：0cdf5b5d0ce1_01.jpg
        self.images = os.listdir(base_dir + 'train')
        # mask：0cdf5b5d0ce1_01_mask.gif
        self.masks = os.listdir(base_dir + 'train_masks')
        self.mode = mode
        # 数据变换
        self.transform = transform
    def __len__(self):
        return len(self.idx_list)
    
    def __getitem__(self, index):
        # 0cdf5b5d0ce1_01.jpg
        image_file = self.images[self.idx_list[index]]
        # 0cdf5b5d0ce1_01.jpg -> 0cdf5b5d0ce1_01 + mask.gif -> 0cdf5b5d0ce1_01_mask.gif
        mask_file = image_file[:-4] + '_mask.gif'
        # 转换成PIL图片
        image = PIL.Image.open(os.path.join(self.base_dir, 'train', image_file))
        if self.mode == 'train':
            mask = PIL.Image.open(os.path.join(self.base_dir, 'train_masks', mask_file))
            if transforms is not None:
                image = self.transform(image)
                mask = self.transform(mask)
                mask[mask!=0] = 1.0
                return image, mask.float()
        else:
            if self.transforms is not None:
                image = self.transform(image)
                return image

初始化及数据集划分

base_dir = './Carvana/'
batch_size = 32
num_workers = 4
# 原图为1918*1280
img_size = (256, 256)
transform = transforms.Compose([
    # 不用ToPILImage了
    transforms.Resize(img_size),
    transforms.ToTensor()
])

# 数据集划分
train_idxs, val_idxs = train_test_split(range(len(os.listdir(base_dir + 'train_masks'))), test_size=0.3)
# 实例化
train_data = CarvanaDataset(base_dir=base_dir, idx_list=train_idxs, transform=transform)
val_data = CarvanaDataset(base_dir=base_dir, idx_list=val_idxs, transform=transform)
# 读入
# 老规矩了
train_loader = DataLoader(dataset=train_data, batch_size=batch_size, num_workers=num_workers, shuffle=True)
val_loader = DataLoader(dataset=val_data, batch_size=batch_size, num_workers=num_workers, shuffle=False)

• 读取一个数据试试

image, mask = next(iter(train_loader))
print(image.shape, mask.shape)
# subplot千万被s大写
plt.subplot(121)
plt.imshow(image[0, 0])
plt.subplot(122)
plt.imshow(mask[0, 0], cmap='gray')

结果如下：

损失函数和优化器

# 使用Binary Cross Entropy Loss
criterion = nn.BCEWithLogitsLoss()
# Adam
optimizer = optim.Adam(unet.parameters(), lr=1e-3, weight_decay=1e-8)
# 多卡并行计算
unet = nn.DataParallel(unet).cuda()

• 也可以使用自定义损失函数

class DiceLoss(nn.Module):
    def __init__(self, weight=None, size_average=True):
        super(DiceLoss, self).__init__()
    
    def forward(self, inputs, targets, smooth=1):
        # 变成(0, 1)
        inputs = torch.sigmoid(inputs)
        inputs = inputs.view(-1)  # 拉平
        targets = targets.view(-1)
        # 1 1才可以
        intersection = (inputs * targets).sum()
        # dice的定义
        dice = (2. * intersection + smooth) / (inputs.sum() + targets.sum() + smooth)

# 测试一下
newcriterion = DiceLoss()
unet.eval()
image, mask = next(iter(val_loader))
out_unet = unet(image.cuda())
loss = newcriterion(out_unet, mask.cuda())
print(loss)

评价指标

# 评价指标，dice co effection
# pred:预测的
# target:真实的
def dice_coeff(pred, target):
    eps = 0.0001
    num = pred.size(0)
    m1 = pred.view(num, -1)  # 拉平
    m2 = target.view(num, -1)
    intersection = (m1 * m2).sum()
    # 这个是dice的公式
    return (2. * intersection + eps) / (m1.sum() + m2.sum() + eps)
    

训练与评估

• 训练

def train(epoch):
    # model.train()的作用是启用 Batch Normalization 和 Dropout
    unet.train()
    train_loss = 0
    for data, mask in train_loader:
        # 转换为GPU的张量
        data, mask = data.cuda(), mask.cuda()
        # 梯度置为0
        optimizer.zero_grad()
        # 训练集上的输出
        output = unet(data)
        # print(output.shape, mask.shape)
        
        # 计算损失
        loss = criterion(output, mask)
        # 将loss反向传播回网络
        loss.backward()
        # 使用优化器更新模型参数
        optimizer.step()
        # data.size(0) data的第1维为为32 一个batch_size为32
        train_loss += loss.item() * data.size(0)  # 每一批样本的损失值之和
    train_loss = train_loss / len(train_loader.dataset)
    print('Epoch:{} \t Training Loss:{:.6f}'.format(epoch, train_loss))

• 验证

def val(epoch):
    print('current learning rate:', optimizer.state_dict()['param_groups'][0]['lr'])
    unet.eval()
    val_loss = 0
    dice_score = 0
    with torch.no_grad():
        for data, mask in val_loader:
            data, mask = data.cuda(), mask.cuda()
            output = unet(data)
            loss = criterion(output, mask)
            val_loss += loss.item() * data.size(0)  # 每一批样本的损失值之和
            dice_score += dice_coeff(torch.sigmoid(output).cpu(), mask.cpu()) * data.size(0)
    val_loss = val_loss / len(val_loader.dataset)  # 算平均的
    dice_score = dice_score / len(val_loader.dataset)  # 算平均的
    print('Epoch:{} \t Validation Loss:{:.6f}, dice score:{:.6f}'.format(epoch, val_loss, dice_score))

• 开始训练

epochs = 100
for epoch in range(1, epochs + 1):
    train(epoch)
    val(epoch)