PyTorch 笔记

本文章归类于 代码基础 类别中，因为 PyTorch 实际上是个代码框架，并不是DL理论内容

本文章中的相关代码已经公开在 GitHub 上，详情参考链接： GitHub 仓库

1 Dataset

获取数据集和对应的 label

有以下主要功能：

1. 如何获取每一个数据和对应的 label
2. 如何获取到数据集的总量大小

常用数据集格式是：一个文件夹放置所有的图像，另外将不同的图像的 label 写在对应的 txt 文件中

1.1 Override

1.1.1 __init__()

一般放全局变量，包括路径等内容，这里使用 hymenoptera_data 数据集作为例子

hymenoptera_data/
└── train/
    ├── ants/
    │   ├── xxx1.jpg
    │   ├── xxx2.png
    │   └── ...
    └── bees/
        ├── xxx1.jpg
        ├── xxx2.png
        └── ...

from torch.utils.data import Dataset
def __init__(self, root: str, label: str):
    """数据的初始化工作

    Args:
        root (str): 数据的根路径
        label (str): 某一个子文件夹的名称，也就是label
    """
    # 这里是根据不同的数据集来更改的
    # 以 hymenoptera_data 数据作为例子， label 下存放了所有的数据图片
    self.root_dir = root # root 路径
    self.label_dir = label # 标签信息
    self.img_path = os.path.join(self.root_dir, self.label_dir) # root 路径地址拼接
    self.img_list = os.listdir(self.img_path)

1.1.2 __getitem__()

获取某一个数据的具体信息

def __getitem__(self, index: int):
    """获取某一个数据的具体信息

    Args:
        index (int): 数据索引号

    Returns:
        img (ndarray, RGB), label (str)
    """

    img_path = os.path.join(self.img_path, self.img_list[index])
    img = cv2.imread(img_path) # OpenCV 默认 BGR
    # 读取失败检查
    if img is None:
        raise FileNotFoundError(f"Failed to read image: {img_path}")
    # 转成 RGB
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    label = self.label_dir

1.1.3 使用实例

if __name__ == '__main__':
    data_root = '/Users/rainer/MyFiles/Code/Learning/Pytorch/dataset/hymenoptera_data/train'

    ant_data = Hymenoptera(root=data_root, label='ants')
    bee_data = Hymenoptera(root=data_root, label='bees')

    train_dataset = ant_data + bee_data
    print(len(ant_data), len(bee_data))
    print(len(train_dataset))

    test_img, test_label = ant_data[3]  # 简写，等同于 __getitem__

    # OpenCV 显示需要 BGR，所以转回来
    img_bgr = cv2.cvtColor(test_img, cv2.COLOR_RGB2BGR)

    cv2.imshow(test_label, img_bgr)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

2 Dataloader

为网络的输入提供不同的数据形式

Dataloader 一般不需要自己重写，只需要了解每一个参数的作用即可

from torch.utils.data import DataLoader
test_loader = DataLoader(
    test_data,
    batch_size=4,    # 每次读取数据数量
    shuffle=True,    # 是否每次打乱顺序
    num_workers=0,    # 子进程个数
    drop_last=False    # 丢弃多余数据
)

3 Tensorboard

Tensorborad 让我们可以方便的查看结果

3.1 常用 API

1. writer.add_image() 用于添加图像，需要指定图片的类型 (cv/numpy)
2. writer.add_scalar() 添加标量信息，常用的 loss 可以用这个数据写

主要有两个轴，纵坐标是参数 scalar_value 横坐标 global_step

3.2 代码示例

from torch.utils.tensorboard import SummaryWriter
import cv2 as cv

if __name__ == '__main__':

    writer = SummaryWriter(log_dir='logs') # 将事件存储到目标logs文件夹
    img_path = 'dataset/hymenoptera_data/train/ants/0013035.jpg'
    img_cv = cv.imread(img_path)
    print(img_cv.shape) # HWC

    # 1. 添加单个图片演示
    writer.add_image('test', img_cv, 0, dataformats='HWC')

    # 2. 添加某个标量测试
    writer.add_scalar(tag='loss', scalar_value=0.5)

    # 3. 添加一连串的数据，类似于绘制函数图像
    for i in range(100):
        writer.add_scalar(tag='loss2', scalar_value=2 * i, global_step=i)

    writer.close()

3.3 查看结果

注意如果 tb 的输出文件都在同一个文件下方，则会重叠

因此实际使用的时候一般区分不同的子文件夹的

 tensorboard --logdir=logs --port xxx

并在对应的端口下方查看结果

4 Transform

Transform 下方是内部放置了很多的工具，不用再造轮子的

使用方法

from torchvision import transforms

4.1 常用 API

1. Tensor 格式转化

# 1.将 PIL 文件转为 Tensor格式
    img_path = '/root/TestPytorch/dataset/test_data/train/ants/0013035.jpg'
    img = Image.open(img_path)
    tensor_trans = transforms.ToTensor()
    tensor_img = tensor_trans(img) #转化格式为tensor
    # print(tensor_img)
    
# 2.使用 opencv
    img_cv = cv.imread(img_path)
    tensor_cv = tensor_trans(img_cv)
    # print(tensor_cv)
    writer = SummaryWriter('/root/TestPytorch/logs')
    writer.add_image('img_cv', tensor_cv, global_step=10)

2. Normalize

# 3.使用 Normalize 归一化库
    # Normalize 参数包括 mean 代表均值，RGB通道则为3个值的数组，std代表标准差，公式如下
    # output[channel] = (input[channel] - mean[channel]) / std[channel]
    print(tensor_cv[0][0][0])
    normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5],std = [0.5, 0.5, 0.5])
    img_normalize = normalize(tensor_cv)
    writer.add_image('img_normalize', img_normalize, global_step=10)
    print(img_normalize[0][0][0])

3. Resize

# 4. 使用 Resize
    # 默认使用 (H, W) 参数进行输入，如果输入单一值，将会按照最小边匹配进行缩放
    resize = transforms.Resize((512, 512))
    print(tensor_cv.size())
    img_resize = resize(tensor_cv)
    print(img_resize.size())
    writer.add_image('img_resize', img_resize, global_step=10)

3. Compose

# 5.使用 Compose 整合操作流程
    trans_compose = transforms.Compose([
    transforms.Resize((512, 512)),transforms.ToTensor(),
    ])
    img_compose = trans_compose(img) # 先进行resize 再进行totensor
    writer.add_image('img_compose', img_compose, global_step=10)

5 Network

这里是构建神经网路的组建，包含了常见的一些 API

实际用的类是 nn.module

5.1 nn.Module

最基本的模块，网络均继承自这个类，使用方法例如

forward 会对输入依次做出内部的操作

import torch
import torch.nn as nn
class MyModel(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
    def forward(self, x):
        output = x + 1
        return output

if __name__ == '__main__':
    model = MyModel()
    x = torch.tensor(1.0)
    output = model(x)
    print(output)

5.2 Conv

卷积操作，常见参数，需要注意输入 tensor 的格式要求

torch.nn.functional.conv2d 要求输入张量 input 的形状严格为 (N, C_in, H, W)

• N：batch size
• C_in：输入通道数
• H, W：图像高、宽

如果操作 reshape 的话，不可以全部限制死，通过将 N 指定为 -1 可以让函数自行计算 N 的大小

output = torch.reshape(output,(-1, 3, 30, 30))

# input: 输入tensor参数，(minibatch, in_channel, H, W)
# weight: 卷积核，被卷积的一个矩阵，会反复的移动 (需要有channel参数)
# bias: 偏置值
# stride: 卷积移动的步长（默认为1），也可以是一个元组(a,b)分别控制横向和纵向步长
# padding: 填充数值为0的行列，可以用元组指定，也可以是string的"same"表示控制输出同一个大小，只能在stride = 1时使用
import torch
output = torch.nn.functional.conv2d(input, kernel,stride=(1,1))

一个卷积层可以这样写（完整例子，下方省略类似例子）

class Model(torch.nn.Module):
    def __init__(self):
        # 注意这里初始化的调用方法
        super(Model, self).__init__()
        self.conv1 = torch.nn.Conv2d(in_channels=3, out_channels=6, kernel_size=3, stride=1, padding=0)
    def forward(self, x):
        return self.conv1(x)

5.3 Pool

池化层，也就是取某一个值来代表某组数据，这里以最大池化层作为例子，简略给出例子

# torch.nn.MaxPool2d()
# input 注意应该是四维array
# kernel_size : 池化核大小
# stride : 窗口的步长,默认是池化核的大小
# padding : 周边填充
# dilation     : 卷积中插入的数值，一般不用
# ceil_mode(默认False) : ceil模式的开关，向上取整，如果True则保留外围一圈;floor是向下取整

self.pool1 = torch.nn.MaxPool2d(kernel_size=2)

5.4 Activation

激活函数，这里以 ReLU 作为例子，简略例子

self.nonlinear = torch.nn.ReLU(inplace=False) # 创建个新数组，并非原地处理

5.5 Linear

线性层

torch.nn.Linear(in_features, out_features, bias=True)
# in_feature: 代表输入的数量
# out_feature: 代表输出的数量
# bias: 表示是否启用 偏置

self.linear1 = torch.nn.Linear(in_features=196608, out_features=10)

5.6 Sequential

把一组操作整合在一起

def __init__(self):
    super(Model, self).__init__()
    self.model1 = nn.Sequential(
        nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding="same"),
        nn.MaxPool2d(kernel_size=2),
        nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding="same"),
        nn.MaxPool2d(kernel_size=2),
        nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding="same"),
        nn.MaxPool2d(kernel_size=2),
        nn.Flatten(),
        nn.Linear(in_features=1024, out_features=64),
        nn.Linear(in_features=64, out_features=10)
    )
def forward(self, x):
    x = self.model1(x)
    return x

5.7 Loss

常用的有两种 ，MSE 和 Cross Entrop

inputs = torch.tensor([0, 1, 2],dtype=torch.float32)
targets = torch.tensor([0, 1, 5],dtype=torch.float32)
loss = torch.nn.MSELoss(reduction='mean')
result = loss(inputs, targets)

...
entropy_loss = torch.nn.CrossEntropyLoss()
entropy_result = entropy_loss(x_output, targets)

5.8 Backward

反向传播，根据 loss 和 gradient 进行参数的更新

model = Model()
loss = nn.CrossEntropyLoss()

for data in dataloader:
    inputs, labels = data
    output = model(inputs)
    result_loss = loss(output, labels)
    result_loss.backward() # 反向传播，注意是对结果值进行反向传播，是优化器的使用基础
    break

5.9 Optimizer

每一个优化器都有自己的用法和参数，通用的有 lr

注意优化的时候需要将梯度归零，否则前几轮会遗留影响

model = Model()
loss = nn.CrossEntropyLoss() # 损失函数
optimizer = torch.optim.SGD(model.parameters(), lr=0.001) # 设置优化器

for epoch in range(20):
    # 假设进行20轮训练，则可以观察每一轮后的loss值为多少
    running_loss = 0.0
    for data in dataloader:
        inputs, labels = data
        output = model(inputs)
        result_loss = loss(output, labels)
        optimizer.zero_grad() # 首先清空梯度，避免之前的循环造成影响
        result_loss.backward() # 反向传播，注意是对结果值进行反向传播，是优化器的使用基础，改变了梯度值
        optimizer.step() # 利用梯度进行反向传播
        running_loss += result_loss.item() # 计算这一轮中所有loss的总和
    print(epoch, running_loss)

6 Model

模型主要分为两个部分，模型的保存与读取

# 保存模型
vgg16 =torchvision.models.vgg16()
torch.save(vgg16, "./model/vgg16.pth")

# 读取文件
model = torch.load("./model/vgg16.pth", weights_only=False)
print(model)

或者可以只保存字典形式，这是更加推荐的方法

torch.save(vgg16.state_dict(), "./model/vgg16_dict.pth")
model = torch.load("./model/vgg16_dict.pth")
print(model) # 这是一个字典

model = torchvision.models.vgg16()          # 先实例化网络结构
state_dict = torch.load("./model/vgg16_dict.pth")  # 加载权重字典
model.load_state_dict(state_dict)                       # 把权重载入模型
model.eval()                                            # 推理模式

7 训练模型

以 CIFAR 10 作为例子

pytorch 基本框架已经完毕，下文是实际训练相关的操作

7.1 准备数据集

直接使用官方的已有的方式加载，如果是自己的数据集可能需要额外的客制化

import torch
from torchvision import transforms
train_data = torchvision.datasets.CIFAR10(root='../dataset', train=True, transform=transforms.ToTensor(), download=True)
test_data = torchvision.datasets.CIFAR10(root='../dataset', train=False, transform=transforms.ToTensor(), download=True)

7.2 加载数据

使用 Dataloader 加载，DataLoader 一般来说不用客制化

from torch.utils.data import DataLoader
# 加载数据集
train_data_loader = DataLoader(train_data, batch_size=64)
test_data_loader = DataLoader(test_data, batch_size=64)

7.3 构建模型

# 模型参见上文，此处省略
from .model import NetworkDemo
model = NetworkDemo()

7.4 Loss 和 优化器

loss = nn.CrossEntropyLoss()
learning_rate = 0.01
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

7.5 训练模型

model.train() 对于一些特殊的网络层有作用，例如 Normalize 或者 Drop out

for i in range(epoch):
    print(f"------ Train Epoch [{i + 1:2}] ------")

    # 训练步骤
    model.train()
    for data in train_data_loader:
        imgs, target = data
        outputs = model(imgs)
        loss = loss_func(outputs, target)
        # 优化器
        optimizer.zero_grad() # 务必清空梯度
        loss.backward()
        optimizer.step()

        train_step += 1
        if train_step % 100 == 0:
            print(f"Iteration [{train_step}], loss {loss.item()}")
            writer.add_scalar("train_loss", loss.item(), train_step)

7.6 验证模型

注意要关掉梯度功能， with torch.no_grad():

# 测试功能
model.eval()
test_loss = 0
with torch.no_grad():    # 关闭梯度功能
    for data in test_data_loader:
        imgs, targets = data
        outputs = model(imgs)
        loss = loss_func(outputs, targets)
        test_loss += loss

    test_loss = test_loss / len(test_data_loader)

7.7 使用 TensorBoard 记录数据

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter("../logs")
...

if train_step % 100 == 0:
    print(f"Iteration [{train_step}], loss {loss.item()}")
    writer.add_scalar("train_loss", loss.item(), train_step)
...

writer.add_scalar("test_loss_epoch", test_loss.item(), epoch)

7.8 保存每一轮的训练模型

torch.save(model, f"../train/model/model_{epoch}.pth")

# 或者可以只保存参数
torch.save(model.state_dict(), f"../train/model/model_{epoch}.pth")

7.9 GPU 训练

GPU 训练的本质就是把数据和模型都放在 GPU 上即可

# 选择训练设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = NetworkDemo()
model.to(device) # 模型放到对应设备上

for i in range(epoch):
    print(f"------ Train Epoch [{i + 1:2}] ------")

    model.train()
    for data in train_data_loader:
        imgs, target = data
        imgs = imgs.to(device)    # 数据放到 GPU
        target = target.to(device)    # 数据放到 GPU
        outputs = model(imgs)
        loss = loss_func(outputs, target)

8 模型验证

8.1 读取数据

cv2 的读取相对麻烦一些，可以在实际项目中自行使用（cv2 速度更快）

这里用 PIL 因为更加方便

import cv2
import torchvision
from PIL import Image

image_path = "../dataset/test/plane.jpeg"
img = cv2.imread(image_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# 使用 PIL 读取
img = Image.fromarray(img)

transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((64, 64)),
    torchvision.transforms.ToTensor(),
])

img = transform(img)
print(img.shape)

8.2 加载模型

# 这里是保存时候的代码
torch.save(model.state_dict(), f"../train/model/model_{i}.pth")

# 加载使用字典
model = NetworkDemo()
state_dict = torch.load("../train/model/model_8.pth")
model.load_state_dict(state_dict)
model.eval()

8.3 进行推理

这里有个坑就是，把输入需要 reshape 到和网络模型完全一致

# 构建 batch size
img = torch.reshape(img, (1, 3, 32, 32))
with torch.no_grad():
    output = model(img)
print(output)
print(output.argmax(1))     # 打印最大值

8.4 查看结果

这里用 CIFAR 作为例子，可以在 dataset 中查看每一个 label 的实际代表

# 可以打印数据集中的那个类别
from torchvision.datasets import CIFAR10
dataset = CIFAR10(root="../dataset", train=False, download=False)
print(dataset.classes[output.argmax(1)])

这里我是用了一张飞机的图片