Pytorch简易入门教程

声明：本教程依照Bilibili Up主我是土堆出品系列视频编撰，有兴趣的读者可以去原视频进行详细学习。

1. 环境配置

使用Anaconda作为包管理工具，创建初始环境、安装torch，cuda和cudnn。

推荐安装Jupyter notebook解锁新天地。

使用代码检测环境是否配置成功：

import torch
print("Pytorch version：")
print(torch.__version__)
print("CUDA Version: ")
print(torch.version.cuda)
print("cuDNN version is :")
print(torch.backends.cudnn.version())
print(torch.cuda.is_available())

笔者配置的环境如下：

Pytorch version：
2.5.1
CUDA Version: 
12.4
cuDNN version is :
90100
True

在学习过程中，我们遇到问题往往会有多种信息来源进行参考，那么这时我们就会面临一个信息的筛选。建议在学习时依照官方文档进行对照，常用的函数为dir()以及help()，后者在例如Pycharm、Visuial Studio Code等编辑器中可以用快捷键Ctrl+左键进行替代。

2. PyTorch数据加载初认识

在PyTorch中，读取数据主要涉及两种方式：Dataset以及Dataloader。

对于Dataset，它是对于数据集的抽象，定义了数据的类型、来源和格式，并且进行编号。通常Dataset只处理单条数据。

而对于Dataloader，它是一个数据加载器，用于控制数据的加载方式，主要负责将Dataset中的数据批量（batch）加载出来，同时提供了例如shuffle、batch_size等基本功能，返回一个可迭代对象，每次迭代返回一个batch。

这里我们以一个蚂蚁蜜蜂图片数据集为例（下载链接）进行具体的学习。

数据集的基本结构如下：

hymenoptera_data/
|--train/
	|--ants/
	└--bees/
└--val/
	|--ants/
	└--bees/

其中，ants和bees目录下存放对应的图片。

下面我们来具体构造一个Dataset。构造一个Dataset一般只需要重写torch提供的类的三个函数：__init__，__getitem__和__len__，分别执行初始化、获得数据与标签、获取数据集长度的功能。

对于我们本案例中的数据集，将相同标签的所有数据条目放在同一目录下，因此文件夹名ants和bees为我们的二分类标签。在另一种数据集的组织方式中，每一条数据会有对应的同名txt文件，里面存放相关的标签信息。

# 导入必要的库
from torch.utils.data import Dataset
from PIL import Image
import os

# 创建Dataset类
class MyDataset(Dataset):
    # 初始化方法，接收根目录和标签目录
    # root_dir: 数据集根目录
    # label_dir: 标签目录（子目录名）
    def __init__(self,root_dir,label_dir):
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.path = os.path.join(self.root_dir, self.label_dir)
        self.image_list = os.listdir(self.path)

    # getitem需要获得索引对应的图像和标签
    def __getitem__(self, index):
        image_name = self.image_list[index]
        image_path = os.path.join(self.path, image_name)
        image = Image.open(image_path)
        label = self.label_dir
        return image, label
    
    def __len__(self):
        return len(self.image_list)

# 创建一个测试用实例
root_dir = 'hymenoptera_data/train'
label_dir = 'ants'
ants_dataset = MyDataset(root_dir, label_dir)
image, label = ants_dataset[1]
print(f"Image: {image}, Label: {label}")
image.show()  

root_dir = 'hymenoptera_data/train'
label_dir = 'bees'
bees_dataset = MyDataset(root_dir, label_dir)
image, label = bees_dataset[5]
print(f"Image: {image}, Label: {label}")
image.show()

train_dataset = ants_dataset + bees_dataset

这样，我们就完成了一个基本数据集的加载。总而言之，加载数据集最为核心的部分便是分辨不同数据的格式（图片、语音、文本等），利用不同的工具进行加载，并同时加载对应的label信息。

3. PyTorch图片处理和可视化

在训练过程中，对于时间成本较高的训练过程，我们往往需要实时跟踪训练过程中一些变量的变化来判断训练是否正常；对于训练结果，我们也需要可视化的展示。另一方面，对于一些图片数据，我们需要进行预处理。在这里介绍Tensorboard的使用和torchvision.transform的使用。

TensorBoard 是 TensorFlow 的可视化工具，PyTorch 提供了对它的支持。可以用来可视化：

• 损失函数与准确率变化曲线
• 模型结构（graph）
• 权重直方图
• 图片、音频、文本等
• Embedding 映射

在使用的过程中，Tensorboard会维护一个log目录，下面存放每一次迭代产生的记录。

3.1 用add_scalar方法添加标量数据

函数原型：add_scalar(tag, scalar_value, global_step=None, walltime=None)

from torch.utils.tensorboard import SummaryWriter
from PIL import Image

# 使用add_scalar方法添加标量数据
# 例如，记录一个简单的线性关系 y = 2x
# 这里我们将x作为step，y作为scalar值
# 这将生成100个点，x从0到99，y从0到198
# 即 y = 2 * x
writer = SummaryWriter("logs")
for i in range(100):
    writer.add_scalar("y=2x", 2*i, i)
writer.close()

3.2 用add_image方法添加图片

函数原型：add_image(tag, img_tensor, global_step=None, walltime=None, dataformats='CHW')

参数名	类型	说明
tag	str	标签名，如 "Input/Image"，用于标识图像
img_tensor	Tensor 或 numpy.ndarray	图像数据，推荐格式为 Tensor（如 torch.Tensor）
global_step	int	步数或 epoch（可选）
walltime	float	时间戳（可选）
dataformats	str	图像格式，常用有：'CHW'（默认），'HWC'，'NCHW'（多图像）等

可以看到，使用add_image方法接受的数据为Tensor或者numpy.array类型。而我们通过：

from PIL import Image
image_path = "Pics/01.jpg"
image = Image.open(image_path)

这种方式获得的image的数据格式为PIL.JpegImagePlugin.JpegImageFile,是无法直接作为该方法的参数传递的。因此我们需要进行类型转换。这里就引出了我们transform函数的使用。

3.2.1 转换为numpy.array

利用opencv库的cv2.imread(image_path)或者numpy库的np.array(pil_image)都能够实现将图片读取为numpy.array类型。

import numpy as np
import cv2
from PIL import Image

image_path = "Pics/01.jpg"
image = Image.open(image_path)
print(type(image))

image_np_array = np.array(image)
image_cv_imread = cv2.imread(image_path)
print(type(image_np_array), image_np_array.shape)
print(type(image_cv_imread), image_cv_imread.shape)

<class 'PIL.JpegImagePlugin.JpegImageFile'>
<class 'numpy.ndarray'> (512, 768, 3)
<class 'numpy.ndarray'> (512, 768, 3)

这样就完成了转换，并且注意输出格式为HWC，详见shape内容。之后再利用add_image方法便可以实现在tensorboard中插入图片

# 使用add_image方法添加图片
# 注意：图片需要是numpy数组格式，且通道顺序为HWC（高度, 宽度, 通道）
# 如果你的图片是RGB格式，确保它是numpy数组且通道顺序正确
writer = SummaryWriter("logs")
writer.add_image("test",image_np_array, 0,dataformats='HWC')
writer.add_image("test",image_cv_imread, 1,dataformats='HWC')
writer.close()

3.2.2 使用torchvison.transforms转换为Tensor

利用torchvision库的transforms可以实现将图片转换为多种格式的Tensor。

原理是利用transforms下面的几个不同类实现。

transforms.ToTensor(image):

功能	说明
格式转换	将图像从 PIL.Image或 numpy.ndarray转换为 torch.Tensor
维度调整	图像从形状 (H, W, C)（常见于 PIL/NumPy）转为 (C, H, W)，PyTorch标准格式
数据归一化	将像素值从 [0, 255]的整数范围 → [0.0, 1.0]的浮点数（除以255）

tensor_trans = transforms.ToTensor()
img_tensor = tensor_trans(image)
print(type(img_tensor), img_tensor.shape)

<class 'torch.Tensor'> torch.Size([3, 512, 768])

transforms.Normalize(mean, std):

对每个通道c，将像素值x执行下面的变换：
$$
x’ = \frac{x- \rm mean [c]}{\rm std[c]}
$$

trans_normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
img_normalized = trans_normalize(img_tensor)

transforms.Resize(size):

对图片进行缩放，有两种方式：

transforms.Resize((512, 512))  # 高512、宽512
transforms.Resize(256)  # 最短边为 256，长边等比例缩放