在服务器上使用 GPU 进行深度学习训练时，第一步通常不是直接跑模型，而是先确认环境是否可用。比如 PyTorch 是否正确安装、CUDA 是否可用、GPU 是否被识别、训练流程是否能正常跑通。
下面这段代码就是一个适合入门练习的 PyTorch 示例。它会先检查 A800 GPU 环境，再构造一个简单的合成数据集，然后训练一个全连接分类模型，最后保存模型权重。
完整代码
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
# 1. 检查 CUDA 环境
print("PyTorch version:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
print("CUDA device count:", torch.cuda.device_count())
if torch.cuda.is_available():
print("Current device:", torch.cuda.get_device_name(0))
device = torch.device("cuda:0")
else:
device = torch.device("cpu")
print("Using device:", device)
# 2. 构造一个简单合成数据集
num_samples =10000
input_dim =128
output_dim =10
X = torch.randn(num_samples, input_dim)
y = torch.randint(0, output_dim, (num_samples,))
dataset = TensorDataset(X, y)
dataloader = DataLoader(dataset, batch_size=256, shuffle=True)
# 3. 定义一个简单分类模型
class SimpleClassifier(nn.Module):
def__init__(self, input_dim, hidden_dim, output_dim):
super().__init__()
self.net = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, output_dim),
)
def forward(self, x):
returnself.net(x)
model = SimpleClassifier(input_dim, 256, output_dim).to(device)
# 4. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# 5. 简单训练循环
epochs =5
for epoch inrange(1, epochs +1):
running_loss =0.0
correct =0
total =0
for batch_x, batch_y in dataloader:
batch_x, batch_y = batch_x.to(device), batch_y.to(device)
optimizer.zero_grad()
outputs = model(batch_x)
loss = criterion(outputs, batch_y)
loss.backward()
optimizer.step()
running_loss += loss.item() * batch_x.size(0)
_, predicted = outputs.max(1)
total += batch_y.size(0)
correct += predicted.eq(batch_y).sum().item()
epoch_loss = running_loss / total
epoch_acc = correct / total
print(f"Epoch {epoch}/{epochs} - Loss: {epoch_loss:.4f} - Acc: {epoch_acc:.4f}")
# 6. 保存模型
torch.save(model.state_dict(), "simple_classifier.pth")
print("Model saved to simple_classifier.pth")
逐行代码解释
1. 导入依赖库
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
·import torch：导入 PyTorch 主库，用于张量计算、模型构建和 GPU 管理。
·import torch.nn as nn：导入神经网络模块，包含常用层、损失函数和模型基类。
·import torch.optim as optim：导入优化器模块，例如 Adam、SGD 等。
·from torch.utils.data import DataLoader, TensorDataset：导入数据集和数据加载器，用于批量读取数据。
2. 检查 CUDA 环境
print("PyTorch version:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
print("CUDA device count:", torch.cuda.device_count())
·torch.__version__：获取当前 PyTorch 版本号。
·torch.cuda.is_available()：判断 CUDA 是否可用，也就是 PyTorch 是否能使用 GPU。
·torch.cuda.device_count()：返回当前可用的 GPU 数量。
if torch.cuda.is_available():
print("Current device:", torch.cuda.get_device_name(0))
device = torch.device("cuda:0")
else:
device = torch.device("cpu")
·torch.cuda.get_device_name(0)：获取第 0 号 GPU 的名称，例如 NVIDIA A800。
·torch.device("cuda:0")：指定使用第 0 号 GPU。
·torch.device("cpu")：如果 CUDA 不可用，则回退到 CPU。
print("Using device:", device)
·打印最终使用的设备，方便确认模型和数据是否会放到 GPU 上。
3. 构造合成数据集
num_samples =10000
input_dim =128
output_dim =10
·num_samples：样本数量，这里生成 10000 条数据。
·input_dim：输入特征维度，每个样本有 128 维特征。
·output_dim：分类类别数，这里是 10 类。
X = torch.randn(num_samples, input_dim)
y = torch.randint(0, output_dim, (num_samples,))
·torch.randn(num_samples, input_dim)：生成服从标准正态分布的随机张量，形状为 [10000, 128]。
·torch.randint(0, output_dim, (num_samples,))：生成 0 到 9 之间的随机整数标签，形状为 [10000]。
这里使用的是随机数据，所以标签本身没有真实语义，主要用于验证训练流程是否跑通。
dataset = TensorDataset(X, y)
dataloader = DataLoader(dataset, batch_size=256, shuffle=True)
·TensorDataset(X, y)：将特征 X 和标签 y 打包成一个数据集。
·DataLoader(...)：创建数据加载器。
·batch_size=256：每个批次包含 256 个样本。
·shuffle=True：每个 epoch 打乱数据顺序，有助于训练稳定性。
4. 定义分类模型
class SimpleClassifier(nn.Module):
def__init__(self, input_dim, hidden_dim, output_dim):
super().__init__()
self.net = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, output_dim),
)
·class SimpleClassifier(nn.Module)：定义一个继承自 nn.Module 的模型类。
·super().__init__()：调用父类初始化方法，这是 PyTorch 模型的标准写法。
·nn.Sequential(...)：按顺序堆叠网络层。
·nn.Linear(input_dim, hidden_dim)：全连接层，将 128 维输入映射到隐藏层维度。
·nn.ReLU()：激活函数，增加非线性表达能力。
·nn.Linear(hidden_dim, output_dim)：最后一层输出 10 维 logits，对应 10 个类别。
def forward(self, x):
returnself.net(x)
·forward 是模型前向传播方法，输入 x 后返回网络输出。
model = SimpleClassifier(input_dim, 256, output_dim).to(device)
·创建模型实例，隐藏层维度为 256。
·.to(device) 将模型参数移动到 GPU 或 CPU。
5. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
·使用交叉熵损失，适用于多分类任务。
·CrossEntropyLoss 内部已经包含 Softmax，所以模型最后一层不需要再加 Softmax。
optimizer = optim.Adam(model.parameters(), lr=1e-3)
·使用 Adam 优化器。
·model.parameters() 表示优化模型中所有可学习参数。
·lr=1e-3 设置学习率为 0.001。
6. 训练循环
epochs =5
·训练 5 个 epoch，即完整遍历数据集 5 次。
for epoch inrange(1, epochs +1):
running_loss =0.0
correct =0
total =0
·running_loss：累计当前 epoch 的损失。
·correct：累计预测正确的样本数。
·total：累计样本总数。
for batch_x, batch_y in dataloader:
batch_x, batch_y = batch_x.to(device), batch_y.to(device)
·从数据加载器中取出一个批次的数据。
·将数据和标签移动到同一设备，否则会报设备不一致错误。
optimizer.zero_grad()
·清空上一批次的梯度。
·如果不清零，梯度会累加，可能导致训练异常。
outputs = model(batch_x)
·前向传播，得到模型输出。
loss = criterion(outputs, batch_y)
·计算当前批次的损失。
loss.backward()
·反向传播，计算每个参数的梯度。
optimizer.step()
·根据梯度更新模型参数。
running_loss += loss.item() * batch_x.size(0)
·loss.item() 获取当前批次平均损失。
·乘以 batch_x.size(0) 得到当前批次总损失，便于后续求整个 epoch 的平均损失。
_, predicted = outputs.max(1)
·在类别维度上取最大值，得到预测类别。
·_ 表示不需要最大值本身，只需要索引。
total += batch_y.size(0)
correct += predicted.eq(batch_y).sum().item()
·累加当前批次样本数。
·predicted.eq(batch_y) 判断预测是否等于真实标签。
·.sum().item() 统计正确预测数量。
epoch_loss = running_loss / total
epoch_acc = correct / total
print(f"Epoch {epoch}/{epochs} - Loss: {epoch_loss:.4f} - Acc: {epoch_acc:.4f}")
·计算当前 epoch 的平均损失和准确率。
·打印训练进度。
由于标签是随机生成的，准确率不会很高，10 分类任务下接近 0.1 属于正常现象。这里重点是验证训练流程，而不是追求高准确率。
7. 保存模型
torch.save(model.state_dict(), "simple_classifier.pth")
print("Model saved to simple_classifier.pth")
·model.state_dict()：获取模型参数字典。
·torch.save(...)：将参数保存到文件。
·这种方式保存的是模型权重，加载时需要先重建模型结构。
运行方式
在服务器上激活对应 conda 环境后，可以直接运行：
python train_simple.py
如果希望指定某张 GPU 运行，可以使用：
CUDA_VISIBLE_DEVICES=0 python train_simple.py
多卡训练提示
当前服务器有 8 张 A800，如果后续想利用多卡，可以在这段代码基础上加入 nn.DataParallel：
if torch.cuda.device_count() >1:
model = nn.DataParallel(model)
不过正式做分布式训练时，更推荐使用 DistributedDataParallel，它对多卡训练的支持更稳定，也更适合大规模训练场景。