入门Unet：自己总结方便后面回看

入门Unet网络学习

第一部分：网络的结构

Unet_part.py：

目的：实现Unet网络中需要用到的几个基础类：
大部分的解释在代码中：

import torch
import torch.nn.functional as F


class DoubleConv(torch.nn.Module):
    def __init__(self, in_ch, out_ch):
        super(DoubleConv, self).__init__()
        self.doubleconv = torch.nn.Sequential(
            torch.nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
            torch.nn.BatchNorm2d(out_ch),  # 归一化
            torch.nn.ReLU(inplace=True),
            torch.nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
            torch.nn.BatchNorm2d(out_ch),  # 归一化
            torch.nn.ReLU(inplace=True)

        )

    def forward(self, x):
        x = self.doubleconv(x)
        return x


class Down(torch.nn.Module):
    def __init__(self, in_ch, out_ch):
        super(Down, self).__init__()
        self.down = torch.nn.Sequential(
            torch.nn.MaxPool2d(2),
            DoubleConv(in_ch, out_ch)
        )

    def forward(self, x):
        x = self.down(x)
        return x


class Up(torch.nn.Module):
    # 构造函数里面输入的是事先定义的量
    def __init__(self, inputs, outputs, bilinear=True):
        super(Up, self).__init__()
        # 上采样的一种方法：双线性采样
        if bilinear:
            self.up = torch.nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        else:
            self.up = torch.nn.ConvTranspose2d(inputs // 2, outputs // 2, 2, stride=2)
        self.dconv = DoubleConv(inputs, outputs)

    # forward（*x）里面x是传入的tensor
    def forward(self, x, y):  # forward(*input)是通过nn.Module 的__call__方法调用的,就相当于调用了模型就是直接调用它的forward函数，y=model(x)，这个x就是直接传入到forward函数的x参数
        x = self.up(x)
        # Upsample-62        [-1, 128, 256, 256] 
        # 说明上采样后照片像素大小一样，但是为了确保下面需要进行对其
        diffx = torch.tensor([y.size()[2] - x.size()[2]])#其实可以不用tensor
        diffy = torch.tensor([y.size()[3] - x.size()[3]])
        x = F.pad(x, [diffx // 2, diffx - diffx // 2,
                      diffy // 2, diffy - diffy // 2])
        # pad的用法：对x进行0（默认）扩充，后面的列表表示从x的最后一个维度开始，列表的后面两个数分别表示对这个维度两头扩充的维数
        # 负数也有效不过会减小维度，所以这里明显是y大于x，（也就是说进过upsample的图片大小往往不能恢复到原大小）
        # " / "表示浮点数除法，返回浮点float结果;" // "表示整数除法
        x = torch.cat([y, x], dim=1)#通道数的合并,而且y和x必须是tensor
        #print(x.size())
        '''torch.Size([2, 1024, 64, 64])
        torch.Size([2, 512, 128, 128])
        torch.Size([2, 256, 256, 256])
        torch.Size([2, 128, 512, 512])'''
        x = self.dconv(x)
        return x


class OutConv(torch.nn.Module):
    def __init__(self, in_ch, out_ch):
        super(OutConv, self).__init__()
        self.outc = torch.nn.Conv2d(in_ch, out_ch, kernel_size=1)

    def forward(self, x):
        x = self.outc(x)
        return x

---

小问题：

1：对于tensor量，x.size()和x.shape返回的是一样
2：关于F.pad中参数列表中数据是int就行，意思就是diffx是int就行（test过没有差别）

---

Unet.py：

from Unet_parts import *
from torchsummary import summary  # 打印torch每层的结构参数


class Unet(torch.nn.Module):
    def __init__(self, in_channels, out_class):
        super(Unet, self).__init__()
        self.n_channels = in_channels#构造函数实现形参的传值
        self.n_classes = out_class

        self.dcon = DoubleConv(in_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        self.down4 = Down(512, 512)
        #输入的size如下，因为进行了dim=1的cat 相当于通道数合并
        # torch.Size([1, 1024, 64, 64]) torch.Size([1, 512, 128, 128])
        # torch.Size([1, 256, 256, 256])torch.Size([1, 128, 512, 512])
        self.up1 = Up(1024, 256)
        self.up2 = Up(512, 128)
        self.up3 = Up(256, 64)
        self.up4 = Up(128, 64)
        self.oconv = OutConv(64, out_class)

    def forward(self, x):
        x1 = self.dcon(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)
        x = self.oconv(x)
        return x


if __name__ == '__main__':
    unet = Unet(1, 1)
    summary(unet, input_size=(1, 512, 512), device='cpu')
    #这里的效果是可以方便看到整个网络情况

summary效果图片：

这里-1,意思是初始的in_channels还未赋值，这里只是检查网络的结构

第二部分：数据集的构成

dataset.py

from torch.utils.data import Dataset, DataLoader
import glob  
import os
import cv2
import random


class Dataset_Loader(Dataset):
    def __init__(self, data_path):
        # 初始化函数，读取所有data_path下的图片
        self.data_path = data_path
        self.imgs_path = glob.glob(os.path.join(data_path, 'image/*.png'))  # os用来合并路径
        # glob.glob()返回所有匹配的文件路径列表。

    def augment(self, image, flipCode):
        # 使用flip进行数据增强， flipCode为1水平翻转，0为垂直翻转，-1水平加垂直
        flip = cv2.flip(image, flipCode)
        return flip

    def __getitem__(self, index):
        # 根据index读取图片,根据图片路径读取具体图片数据内容
        image_path = self.imgs_path[index]
        # 获得label_path
        label_path = image_path.replace('image', 'label')
        # 读取图片和标签
        image = cv2.imread(image_path)  # 输入的需要的是带索引的路径
        label = cv2.imread(label_path)
        # 输出为(512, 512, 3)

        # 将数据转为单通道的图片
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        label = cv2.cvtColor(label, cv2.COLOR_BGR2GRAY)
        # 输出为(512, 512) 返回值：灰度图片的数组
        image = image.reshape(1, image.shape[0], image.shape[1])
        label = label.reshape(1, label.shape[0], label.shape[1])
        # (1, 512, 512)
        # 注意这里的输出都是数组类型还不是tensor
        #print(image.shape)
        # 处理标签，将像素值为255的改为1
        if label.max() > 1:
            label = label / 255
        # 随机进行数据增强，为2的不做处理
        flipcode = random.choice([-1, 0, 1, 2])
        if flipcode != 2:
            image = self.augment(image, flipcode)
            label = self.augment(label, flipcode)

        return image, label
        # 返回的是image的（1，1,255,255）的数组和label同样大小但进过等比缩小的数组

    def __len__(self):
        # 返回训练集大小
        return len(self.imgs_path)


if __name__ == "__main__":
    dataset_loader = Dataset_Loader("data/train/")  # 路径名不能有中文啊
    print("num:", len(dataset_loader))
    train_loader = DataLoader(dataset=dataset_loader, batch_size=2, shuffle=True)
    for image,label in train_loader:
        print(image.shape)
    # DataLoader返回的image是torch.Size([2, 1, 512, 512])
    # 而函数getitem中的image是(1, 512, 512)和(1, 512, 512)
    # 所以初始的数据集输入tensor第一位是batch_size，第二位是图像通道大小也是net中的in_channels（这里是灰度图）
    # 并且image也转为tensor了
   

这里对所有的输入输出都有注释，可以更清楚的了解数据在其中的处理

问题

1.对于图片进行数据增强的原理
2.还有一些cv中对于图像处理的常用操作

第三部分：训练和测试

train.py

from Unet import Unet
from dataset import Dataset_Loader
from torch import optim
import torch.nn as nn
import torch


def train_net(net, device, data_path, epochs=40, batch_size=1, lr=1e-5):
    # 加载训练集
    dataset_loader = Dataset_Loader(data_path)
    train_loader = torch.utils.data.DataLoader(dataset=dataset_loader, batch_size=batch_size, shuffle=True)

    # 定义优化方法
    optimizer = optim.RMSprop(net.parameters(), lr=lr, weight_decay=1e-8, momentum=0.5)
    criterion = nn.BCEWithLogitsLoss()

    # best_loss初始化为正无穷
    best_loss = float('inf')

    # 训练epoch次
    for epochs in range(epochs):

        for image, label in train_loader:
            #这里因为一共有30张图片，而且batch_size=1，所以train_loader中有30次训练
            optimizer.zero_grad()
            # 将数据拷贝到device中去
            image = image.to(device=device, dtype=torch.float32)
            label = label.to(device=device, dtype=torch.float32)

            pred = net(image)
            # print(pred.size()) torch.Size([1, 1, 512, 512])

            loss = criterion(pred, label)  # 返回的是什么
            # print(loss) tensor(0.6900, device='cuda:0', grad_fn=<BinaryCrossEntropyWithLogitsBackward>)
            print('Loss/train', epochs, loss.item())

            if loss < best_loss:
                best_loss = loss
                torch.save(net.state_dict(), 'best_model.path')
            # 虽然搞不懂，但是5.4和tensor([3])就是可以比较大小的，赋值过程中就进行了类型的转换了
            # 然后保存最优值

            loss.backward()
            optimizer.step()


if __name__ == "__main__":
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    net = Unet(1, 1)

    net.to(device)

    data_path = "data/train"
    train_net(net, device, data_path)
    # 这里都是进行一些调试

问题

1.关于torch.save：

2./.item()将Tensor变量转换为python标量（int float等），其中t是一个Tensor变量，只能是标量，转换后dtype与Tensor的dtype一致。

test.py

从test这里把预测的权重带入原模型生成预测图片

import glob
import numpy as np
import torch
import os
import cv2
from Unet import Unet

if __name__ == "__main__":
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    net = Unet(1, 1)  # 单通道图片且分类为1
    net.to(device)
    # 加载模型参数
    net.load_state_dict(torch.load('best_model.path', map_location=device))
    # 测试模型
    net.eval()
    # 读取所有图片路径
    test_path = glob.glob('data/test/*.png')

    for test_path in test_path:
        save_test_path = test_path.split('.')[0] + '_test.png'
        # split()：拆分字符串。通过指定分隔符对字符串进行切片，并返回分割后的字符串列表。split:['data/test/0','png']
        # 这里就是以'.'为分割符，取其第0维,返回save_test_path：data/test/0.png'
        # os.path.split()：将文件名和路径分割开。

        img = cv2.imread(test_path)  # (512, 512, 3) 输入的需要的是带索引的路径
        # 转为灰度图
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # (512, 512) 因为是黑白图片所以直接把rgb三值都一样，直接合并
        # 转为batch为1，通道为1，大小为512*512的数组
        img = img.reshape(1, 1, img.shape[0], img.shape[1])  # (1, 1, 512, 512)
        # 这里与dataset不同是因为这里没有经过DataLoader但是放入net的size需要一样

        img_tensor = torch.from_numpy(img)  # 数组与tensor([1, 1, 512, 512])的转换

        img_tensor = img_tensor.to(device=device, dtype=torch.float32)

        pred = net(img_tensor)
        # print(pred.shape)   torch.Size([1, 1, 512, 512])

        # print(pred.data.cpu()[0])  tensor([[[0, 1],
        #                                     [2, 3]]], dtype=torch.int32)  这里的0123是随便写的里面其实是512*512的矩阵
        # pred.data.cpu()[0]是torch.Size([1, 512, 512])
        pred = np.array(pred.data.cpu()[0])[0]
        # print(pred.shape) (512, 512) 并且转换为数组
        # pred.data.cpu()[0]表示读入cpu中

        pred[pred >= 0.5] = 225
        pred[pred < 0.5] = 0

        cv2.imwrite(save_test_path, pred)
        # cv2.imwrite(“图片要保存的相对路径”，os.path.splitext(filenames)[0]+os.splittext(filenames)[1],image)
        # 相对路径的只能以cv2.imwrite()所在当前文件夹开头，这里就是data

问题

.item()

将Tensor变量转换为python标量（int float等），其中t是一个Tensor变量，只能是标量，转换后dtype与Tensor的dtype一致。。

.cpu

将数据的处理设备从其他设备（如.cuda()拿到cpu上），不会改变变量类型，转换后仍然是Tensor变量。

.data

与item()相似 ，返回的是tensor

a = np.array(a.data.cpu()[0])[0]

1.首先a是一个放在GPU上的Variable，a.data是把Variable里的tensor取出来，可以看出与a的差别是：缺少了第一行(Variable containing)

2.a.cpu()和a.data.cpu()是分别把a和a.data放在cpu上，其他的没区别，另外：a.data.cpu()，a.cpu().data一样（这里的a必须是tensor）

3.a.data[0] | a.cpu().data[0] | a.data.cpu()[0]是一样的，都是取出第0维的tensor

4.连续取两次第0维就得到了一个512*512矩阵

数组的骚操作

CV2的操作

参考：https://blog.csdn.net/fu6543210/article/details/80835280

初步学习数据的可视化（tensorboardx）

from tensorboardX import SummaryWriter
################################
writer = SummaryWriter('data/log')
#################################
 # tensorboard --logdir=E:\code\Pytorch-Unet\data\log
 writer.add_scalar('Loss', loss.item(), i+epochs*30)
 writer.close()