Lightning基础训练尝试实例

一、训练任务概述

动机：由于后续的课题中会用到类似图像去噪的算法，考虑先用U-Net，这里做一个前置的尝试。

训练任务：分割出图像中的细胞。

数据集：可私

数据集结构：

二、具体实现

U-Net的网络实现是现成的，只需要在网上找一个比较漂亮的实现（一般都是模块化，写的很漂亮）copy就可以了，需要特别注意的是最后整合的模型：

2.1 基础模型模块实现

双卷积模块

class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels, mid_channels=None):
        super().__init__()
        if not mid_channels:
            mid_channels = out_channels
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)

上采样模块

class Up(nn.Module):
    """Upscaling then double conv"""

    def __init__(self, in_channels, out_channels, bilinear=True):
        super().__init__()

        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
        else:
            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
            self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        # input is CHW
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]

        x1 = torch.nn.functional.pad(x1, [diffX // 2, diffX - diffX // 2,diffY // 2, diffY - diffY // 2])

        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)

下采样模块

class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)

输出层

class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)

2.2 整合模块->模型

class UNet(L.LightningModule):
    def __init__(self, n_channels, n_classes, bilinear=False):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = (DoubleConv(n_channels, 64))
        self.down1 = (Down(64, 128))
        self.down2 = (Down(128, 256))
        self.down3 = (Down(256, 512))
        factor = 2 if bilinear else 1
        self.down4 = (Down(512, 1024 // factor))
        self.up1 = (Up(1024, 512 // factor, bilinear))
        self.up2 = (Up(512, 256 // factor, bilinear))
        self.up3 = (Up(256, 128 // factor, bilinear))
        self.up4 = (Up(128, 64, bilinear))
        self.outc = (OutConv(64, n_classes))
    
    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits

    # 对应的层设置检查点，节省显存m，可用可不用
    def use_checkpointing(self):
        self.inc = torch.utils.checkpoint(self.inc)
        self.down1 = torch.utils.checkpoint(self.down1)
        self.down2 = torch.utils.checkpoint(self.down2)
        self.down3 = torch.utils.checkpoint(self.down3)
        self.down4 = torch.utils.checkpoint(self.down4)
        self.up1 = torch.utils.checkpoint(self.up1)
        self.up2 = torch.utils.checkpoint(self.up2)
        self.up3 = torch.utils.checkpoint(self.up3)
        self.up4 = torch.utils.checkpoint(self.up4)
        self.outc = torch.utils.checkpoint(self.outc)
    
    # 定义优化器
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(),lr=0.001)
        return optimizer

    # 定义train的单步流程
    def training_step(self,train_batch,batch_index):
        image,label = train_batch
        image_hat = self.forward(image)
        # U-Net的loss
        loss = nn.functional.mse_loss(image_hat,label)
        return loss
    
    # 定义val的单步流程
    def validation_step(self, val_batch,batch_index):
        image,label = val_batch
        image_hat = self.forward(image)
        # U-Net的loss
        loss = nn.functional.mse_loss(image_hat,label)
        self.log('val_loss',loss)
        return loss

注意：模块可以不需要继承自L.LightningModule，只要最后整合的时候继承自L.LightningModule就可以了。

2.3 数据划分

重定义Dataset类，供数据集划分函数调用，二者要相互配合

class UDataset(Dataset):
    def __init__(self,image_dir,mask_dir,transform=None):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        if transform is not None:
            self.transform = transform
        else:
            self.transform = None
    
    def __getitem__(self, index):
        image = Image.open(self.image_dir[index]).convert('RGB')
        label = Image.open(self.mask_dir[index]).convert('RGB')
        if self.transform is not None:
            image = self.transform(image)
            label = self.transform(label)
        return image,label

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

 定义数据集划分函数（包括"找出文件列表"、"定义数据预处理方式"、“定义批量大小”）

train_image_dir = "./data/train/image/*.png"
train_label_dir = "./data/train/label/*.png"
val_image_dir = "./data/val/image/*.png"
val_label_dir = "./data/val/label/*.png"  

def data_process(train_image_dir,train_label_dir,val_image_dir,val_label_dir):
    # 查找路径下的所有文件，返回文件路径列表
    train_image_list = glob.glob(train_image_dir)
    train_label_list = glob.glob(train_label_dir)
    val_image_list = glob.glob(val_image_dir)
    val_label_list = glob.glob(val_label_dir)

    # 数据处理
    train_data_transform = transforms.Compose([
        transforms.Resize((256,256)),
        transforms.ToTensor()
    ])
    val_data_transform = transforms.Compose([  
        transforms.Resize((256,256)),
        transforms.ToTensor()
    ])

    train_dataloader = data.DataLoader(UDataset(train_image_list,train_label_list,train_data_transform),batch_size=5,shuffle=True)
    val_dataloader = data.DataLoader(UDataset(val_image_list,val_label_list,val_data_transform),batch_size=5,shuffle=False)
    
    return train_dataloader,val_dataloader

2.4 模型验证

在训练之前，要看一下模型的结构有没有错误，用summary打印出网络的结构

    # 模型验证
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = UNet(n_channels=3,n_classes=1).to(device)
    print(summary(model,(3,512,512)))

也可以用其他的方法查看网络结构

2.5 模型训练

加入TensorBoardLogger是为了可视化训练Loss

训练的流程遵循前文的基本流程

    # 创建 TensorBoardLogger
    logger = TensorBoardLogger("tb_logs", name="unet")
    # 创建 Trainer
    trainer = L.Trainer(max_epochs=20, logger=logger)
    # 划分数据集
    train_dataloader,val_dataloader = data_process(train_image_dir,train_label_dir,val_image_dir,val_label_dir)
    # 创建模型
    model = UNet(n_channels=3,n_classes=1)
    # 启动模型训练过程
    trainer.fit(model,train_dataloader,val_dataloader)
    # 保存模型权重
    torch.save(model.state_dict(),'./model.pth')