Pytorch框架学习

Pytorch环境配置

CUDA driver和CUDA runtime版本

  • 确认自己显卡算力 -> 确认自己显卡型号
  • 确认自己可以选择的CUDA Runtime Version
  • 确保自己的CUDA Driver Version>= CUDA Runtime Version
  • 在Pytorch官网上安装选择的版本小于等于nvidia-smi右侧的CUDA version

参考：

参考案例

数据加载和处理

Dataset

PyTorch 的 Dataset 类是一个数据集的抽象类，主要用于定义如何获取数据及其对应的标签。可以用来加载各种格式的数据，例如图片、文本、CSV 文件等。可以继承 torch.utils.data.Dataset 并重写其中的方法来自定义数据集。
对于图片数据，根据数据集存放地址，创建数据集的代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
import os
from torch.utils.data import Dataset
from PIL import Image

# 自定义 Dataset
class ImageDataset(Dataset):
    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_path = os.listdir(self.path) #获取文件夹内所有图片路径

    def __getitem__(self, idx):
        # 根据索引返回数据和标签
        image_name = self.image_path[idx]
        img_item_path = os.path.join(self.root_dir,self.label_dir,image_name) #获取单个图片路径
        img = Image.open(img_item_path)
        label = self.label_dir
        return img,label    

    def __len__(self): #数据集长度
        return len(self.image_path)

# 创建数据集
root_dir = "dataset/train"
ants_label_dir = "ants"
bees_label_dir = "bees"
ants_dataset = myDataset(root_dir,ants_label_dir)
bees_dataset = myDataset(root_dir,bees_label_dir)
# 查看数据样本
print("ants_dataset size:", len(ants_dataset))
print("ants_dataset First sample:", ants_dataset[0])
train_dataset = ants_dataset + bees_dataset

Dataloader

torch.utils.data.DataLoader 是一个数据加载器，用于从 Dataset 中批量加载数据。

例子1：

现在加载一个 CSV 文件，假设它包含两列数据：x 和 y，我们想将它加载为一个 PyTorch 数据集。：

1
2
3
4
5
6
x,y
1,2
2,4
3,6
4,8
5,10
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import pandas as pd
from torch.utils.data import Dataset, DataLoader

# 定义自定义 Dataset
class CSVDataset(Dataset):
    def __init__(self, csv_file):
        # 读取 CSV 文件
        self.data = pd.read_csv(csv_file)
    
    def __len__(self):
        # 数据集的大小
        return len(self.data)
    
    def __getitem__(self, idx):
        # 根据索引返回数据
        x = self.data.iloc[idx, 0]  # 第一列是 x
        y = self.data.iloc[idx, 1]  # 第二列是 y
        return torch.tensor(x, dtype=torch.float32), torch.tensor(y, dtype=torch.float32)

# 创建数据集和 DataLoader
dataset = CSVDataset("data.csv")
dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

# 遍历数据
for batch in dataloader:
    x, y = batch
    print("x:", x)
    print("y:", y)

对于图片数据，PyTorch 提供了 torchvision 工具包，它可以帮助加载常见的图片数据集（如 MNIST、CIFAR-10 等）。以下是一个简单示例：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import torchvision.transforms as transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader

# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),  # 转换为 PyTorch 的张量
    transforms.Normalize((0.5,), (0.5,))  # 标准化到 [-1, 1]
])

# 加载 MNIST 数据集
train_dataset = MNIST(root="./data", train=True, transform=transform, download=True)

# 创建 DataLoader
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# 查看数据
for images, labels in train_loader:
    print("Batch of images shape:", images.shape)  # [64, 1, 28, 28]
    print("Batch of labels shape:", labels.shape)  # [64]
    break

模型构建

nn.Module

nn.Module 是 PyTorch 中所有神经网络模块的基类，它提供了一些基本的方法和属性，例如 forward、parameters、named_parameters 等。

1
2
3
4
5
6
7
8
9
10
11
12
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

    def forward(self, input):
        output = input + 1
        return output

net = Net()
input = torch.tensor(1.0)
output = net(input)
print(output) #输出为tensor(2.)

完整例子

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
# from model import *

#定义训练设备
device = torch.device("cuda:0")

train_data = torchvision.datasets.CIFAR10(root='./dataset', train=True, download=True,
                                          transform=torchvision.transforms.ToTensor())

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



train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练数据集的长度为:{}".format(train_data_size))
print("测试数据集的长度为:{}".format(test_data_size))

# 加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

# 创建网络模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64 * 4 * 4, 64),
            nn.Linear(64, 10)
        )

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

net = Net()
net.to(device)

# 损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn.to(device)
# 优化器
learning_rate = 1e-2
optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate)

# 设置训练网络参数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
# 训练的轮数
epoch = 10

# 添加tensorboard

writer = SummaryWriter(log_dir='./logs_train')
for i in range(epoch):
    print("-------第{}轮训练开始-------".format(i + 1))

    # 训练步骤开始
    net.train(True)
    for data in train_dataloader:
        imgs, targets = data
        imgs = imgs.to(device)
        targets = targets.to(device)
        output = net(imgs)
        loss = loss_fn(output, targets)

        # 优化器优化模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step += 1
        if total_train_step % 100 == 0:
            print("训练次数:{},loss:{}".format(total_train_step, loss.item()))
            writer.add_scalar('train_loss', loss.item(), total_train_step)

    # 测试步骤开始
    net.eval()
    total_test_loss = 0
    total_accuracy = 0
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)
            output = net(imgs)
            loss = loss_fn(output, targets)
            total_test_loss += loss.item()
            accuracy = (output.argmax(1) == targets).sum()
            total_accuracy += accuracy

    print("整体测试集上的Loss：:{}".format(total_test_loss))
    print("整体测试集上的正确率:{}".format(total_accuracy / test_data_size))
    writer.add_scalar('test_loss', total_test_loss, total_test_step)
    writer.add_scalar('test_accuracy', total_accuracy / test_data_size, total_test_step)
    total_test_step += 1

    torch.save(net, "train_{}.pth".format(i))
    print("模型已保存")

writer.close()