centos pytorch如何训练模型

在 CentOS 上用 PyTorch 训练模型，通常分为 环境准备 → 代码编写 → 启动训练 几个步骤。下面给你一个从零开始的实战流程（包含 CPU 和 GPU 情况）。

一、系统与环境准备（CentOS）

1. 更新系统

sudo yum update -y

2. 安装 Python（推荐 3.8+）

sudo yum install -y python3 python3-pip

检查：

python3 --version

二、安装 PyTorch

✅ CPU 版本

pip3 install torch torchvision torchaudio

✅ GPU 版本（CentOS + CUDA）

先确认你有 NVIDIA 显卡：

nvidia-smi

根据 CUDA 版本安装（示例 CUDA 11.8）：

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

验证：

import torch
print(torch.cuda.is_available())

三、准备训练代码（示例）

1. 简单训练示例（MNIST）

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

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

# 数据
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('./data', train=True, download=True,
                   transform=transforms.ToTensor()),
    batch_size=64, shuffle=True)

# 模型
model = nn.Sequential(
    nn.Flatten(),
    nn.Linear(28*28, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
).to(device)

# 损失 & 优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练
for epoch in range(5):
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        loss = criterion(model(x), y)
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

保存模型：

torch.save(model.state_dict(), "model.pt")

四、后台训练（推荐）

1. 使用 nohup

nohup python3 train.py > train.log 2>&1 &

查看日志：

tail -f train.log

2. 使用 screen（推荐）

screen -S torch
python3 train.py
# Ctrl+A D 退出

五、常见问题（CentOS）

❌ 缺少依赖

sudo yum install -y gcc make

❌ pip 版本低

pip3 install --upgrade pip

❌ 显存不足

• 减小 batch_size
• 使用 torch.cuda.empty_cache()

六、进阶建议

• 使用 DistributedDataParallel 多卡训练
• 使用 TensorBoard 可视化
• 使用 conda（比系统 Python 更稳定）

如果你愿意，可以告诉我：

• ✅ CentOS 版本
• ✅ 是否有 GPU
• ✅ 训练什么模型（CNN / Transformer / 自定义）

我可以直接给你 可运行的完整训练脚本。