本文通过一个简单的二维二分类任务，记录使用 PyTorch 完成数据准备、模型创建、训练、验证和模型保存的完整流程。

1. 基础流程

1.1 数据准备

首先随机生成两组 Tensor 数据：将第一组数据的中心移动到 (-1, -1)，将第二组数据的中心移动到 (1, 1)。随后使用 torch.randperm() 打乱样本顺序，并通过 TensorDataset 和 DataLoader 将数据封装成批次，供后续训练和验证使用。

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
# 1. 准备数据
class_0 = torch.randn(500,2) * 0.6
class_0 += torch.tensor([-1.0,-1.0])

class_1 = torch.randn(500,2) * 0.6
class_1 += torch.tensor([1.0,1.0])

feature = torch.cat(
    [class_0,class_1],
    dim=0
)

labels = torch.cat(
    [
        torch.zeros(500,dtype=torch.long),
        torch.ones(500,dtype=torch.long)
    ]
)

indices = torch.randperm(len(feature))

feature = feature[indices]
labels = labels[indices]

train_x = feature[:800]
train_y = labels[:800]

val_x = feature[800:]
val_y = labels[800:]

train_batch = TensorDataset(
    train_x,
    train_y
)

val_batch = TensorDataset(
    val_x,
    val_y
)

train_loader = DataLoader(
    train_batch,
    batch_size=32,
    shuffle=True
)

val_loader = DataLoader(
    val_batch,
    batch_size=64,
    shuffle=False
)

batch_x , batch_y = next(iter(train_loader))
print(batch_x.shape)
print(batch_y.shape)

重点函数：

  • torch.cat([...], dim=...)：沿指定维度拼接多个 Tensor。
  • torch.randperm(n)：返回由 0 到 n - 1 组成的随机排列，可用于打乱样本顺序。
  • TensorDataset(...)：将多个 Tensor 按照第一维一一对应地组合成数据集。
  • DataLoader(..., batch_size=..., shuffle=...)：按批次加载数据；shuffle 用于控制每轮是否重新打乱数据。

1.2 模型创建

下面继承 nn.Module 创建一个全连接神经网络，并将模型移动到 GPU；如果当前环境不支持 CUDA，则自动使用 CPU。

nn.Module 是 PyTorch 中所有神经网络模块的基类，能够注册并管理模型中的子模块与可训练参数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class Classifier(nn.Module):
    def __init__(self):
        super().__init__()

        self.network = nn.Sequential(
            nn.Linear(2,16),
            nn.ReLU(),
            nn.Linear(16, 16),
            nn.ReLU(),
            nn.Linear(16, 2)
        )

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

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Classifier().to(device=device)

重点函数：

  • super().__init__()：调用父类 nn.Module 的构造方法，使子模块和参数能够被正确注册。
  • nn.Sequential(...)：按照定义顺序依次执行其中的网络层。

1.3 开始训练模型

每个批次的训练流程如下：

  1. 清空上一轮累积的梯度。
  2. 前向传播，得到模型输出 logits。
  3. 根据 logits 和真实标签计算损失。
  4. 反向传播，计算各参数的梯度。
  5. 使用优化器更新模型参数。
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
epochs = 30

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(),lr=1e-3)

for epoch in range(epochs):
    model.train()

    some_loss = 0.0
    corrents = 0
    total = 0

    for batch_x , batch_y in train_loader:
        batch_x = batch_x.to(device)
        batch_y = batch_y.to(device)

        optimizer.zero_grad()
        logits = model(batch_x)

        loss = criterion(logits,batch_y)

        loss.backward()
        optimizer.step()

        some_loss += loss.item() * batch_x.size(0)

        predition = torch.argmax(logits,dim=-1)
        corrents += (predition == batch_y).sum().item()
        total += batch_y.size(0)

    total_loss = some_loss / total
    total_accuracy = corrents / total

    print(
        f"epoch={epoch + 1:02d}, "
        f"loss={total_loss:.4f}, "
        f"accuracy={total_accuracy:.2%}"
    )

重点函数：

  • nn.CrossEntropyLoss()：用于多分类任务，内部组合了 LogSoftmax 和 NLLLoss，因此模型应直接输出未经 Softmax 处理的 logits。
  • torch.optim.AdamW(model.parameters(), lr=1e-3)：创建 AdamW 优化器，用于更新模型参数。
  • optimizer.zero_grad()：清空上一轮反向传播得到的梯度。
  • model.train()：将模型切换为训练模式，影响 Dropout、BatchNorm 等特定层的行为。

1.4 使用验证集评估模型

验证阶段不需要计算和更新梯度，因此使用 torch.no_grad() 关闭梯度记录，以减少内存占用并提高计算效率。调用 model.eval() 可以将模型切换为评估模式。

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
model.eval()
eval_loss = 0.0
eval_current = 0
eval_total = 0

with torch.no_grad():
    for batch_x,batch_y in val_loader:
        batch_x = batch_x.to(device)
        batch_y = batch_y.to(device)

        logits = model(batch_x)
        loss = criterion(logits,batch_y)

        eval_loss += loss.item() * batch_x.size(0)
        predition = torch.argmax(logits,dim=-1)

        eval_current += (predition == batch_y).sum().item()

        eval_total += batch_y.size(0)

total_loss = eval_loss / eval_total
total_accuracy = eval_current / eval_total

print("\n验证结果")
print(f"验证loss: {total_loss:.4f}")
print(f"验证准确率: {total_accuracy:.2%}")

重点函数：

  • model.eval()：将模型切换为评估模式。
  • torch.no_grad()：在代码块内关闭梯度记录，不再构建用于反向传播的计算图。

1.5 保存模型

1
2
torch.save(model.state_dict(),"day03_module_basics.pt")
print("\n模型已保存到 day03_module_basics.pt")

重点函数：

  • model.state_dict()：获取模型中所有可学习参数的状态字典。
  • torch.save(..., "name.pt")：将模型参数保存到指定文件中，便于后续加载和推理。

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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
import torch
from torch import nn
from torch.utils.data import TensorDataset,DataLoader


torch.manual_seed(42)
# 1. 准备数据
class_0 = torch.randn(500,2) * 0.6
class_0 += torch.tensor([-1.0,-1.0])

class_1 = torch.randn(500,2) * 0.6
class_1 += torch.tensor([1.0,1.0])

feature = torch.cat(
    [class_0,class_1],
    dim=0
)

labels = torch.cat(
    [
        torch.zeros(500,dtype=torch.long),
        torch.ones(500,dtype=torch.long)
    ]
)

indices = torch.randperm(len(feature))

feature = feature[indices]
labels = labels[indices]

train_x = feature[:800]
train_y = labels[:800]

val_x = feature[800:]
val_y = labels[800:]

train_batch = TensorDataset(
    train_x,
    train_y
)

val_batch = TensorDataset(
    val_x,
    val_y
)

train_loader = DataLoader(
    train_batch,
    batch_size=32,
    shuffle=True
)

val_loader = DataLoader(
    val_batch,
    batch_size=64,
    shuffle=False
)

batch_x , batch_y = next(iter(train_loader))
print(batch_x.shape)
print(batch_y.shape)


# 2. 创建模型
class Classifier(nn.Module):
    def __init__(self):
        super().__init__()

        self.network = nn.Sequential(
            nn.Linear(2,16),
            nn.ReLU(),
            nn.Linear(16, 16),
            nn.ReLU(),
            nn.Linear(16, 2)
        )

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

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Classifier().to(device=device)

# 3. 开始训练
epochs = 30

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(),lr=1e-3)

for epoch in range(epochs):
    model.train()

    some_loss = 0.0
    corrents = 0
    total = 0

    for batch_x , batch_y in train_loader:
        batch_x = batch_x.to(device)
        batch_y = batch_y.to(device)

        optimizer.zero_grad()
        logits = model(batch_x)

        loss = criterion(logits,batch_y)

        loss.backward()
        optimizer.step()

        some_loss += loss.item() * batch_x.size(0)

        predition = torch.argmax(logits,dim=-1)
        corrents += (predition == batch_y).sum().item()
        total += batch_y.size(0)

    total_loss = some_loss / total
    total_accuracy = corrents / total

    print(
        f"epoch={epoch + 1:02d}, "
        f"loss={total_loss:.4f}, "
        f"accuracy={total_accuracy:.2%}"
    )

# 4.测试
model.eval()
eval_loss = 0.0
eval_current = 0
eval_total = 0

with torch.no_grad():
    for batch_x,batch_y in val_loader:
        batch_x = batch_x.to(device)
        batch_y = batch_y.to(device)

        logits = model(batch_x)
        loss = criterion(logits,batch_y)

        eval_loss += loss.item() * batch_x.size(0)
        predition = torch.argmax(logits,dim=-1)

        eval_current += (predition == batch_y).sum().item()

        eval_total += batch_y.size(0)

total_loss = eval_loss / eval_total
total_accuracy = eval_current / eval_total

print("\n验证结果")
print(f"验证loss: {total_loss:.4f}")
print(f"验证准确率: {total_accuracy:.2%}")

# 5. 保存模型
torch.save(model.state_dict(),"day03_module_basics.pt")
print("\n模型已保存到 day03_module_basics.pt")