Linux与PyTorch协同工作流程

一 环境准备与驱动检查

• 更新系统并安装基础工具：sudo apt update && sudo apt install -y python3 python3-pip build-essential dkms linux-headers-$(uname -r)。
• 安装并验证 NVIDIA 驱动：建议通过官方仓库或 Ubuntu 附加驱动安装与你 GPU 匹配的驱动版本，安装后用 nvidia-smi 查看驱动与 CUDA 运行时版本；如需手动安装，可 sudo apt install nvidia-driver-。
• 规划 CUDA/cuDNN 与 PyTorch 的版本匹配：PyTorch 既可通过 pip 使用预编译的 CUDA 版本（如 cu117、cu118），也可通过 conda 安装并携带 cudatoolkit（如 cudatoolkit=11.8）；两者需版本对应，避免不兼容。

二 创建隔离环境并安装PyTorch

• 使用 venv 创建虚拟环境：
  python3 -m venv ~/venvs/pt_env && source ~/venvs/pt_env/bin/activate
• 使用 conda 创建虚拟环境：
  conda create -n pt_env python=3.10 -y && conda activate pt_env
• 安装 PyTorch（按是否有 NVIDIA GPU 选择其一）：
  • CPU 版：pip install torch torchvision torchaudio
  • GPU 版（pip，示例为 CUDA 11.8）：pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
  • GPU 版（conda，示例为 cudatoolkit=11.8）：conda install pytorch torchvision torchaudio cudatoolkit=11.8 -c pytorch -c nvidia
• 验证安装：
  python - <<‘PY’ import torch print(“torch:”, torch.version, “cuda:”, torch.version.cuda) print(“cuda available:”, torch.cuda.is_available()) print(“device count:”, torch.cuda.device_count(), “name:”, torch.cuda.get_device_name(0) if torch.cuda.is_available() else “N/A”) PY

三 训练与推理的标准工作流

• 数据加载与增强：
  from torch.utils.data import DataLoader, random_split
  from torchvision import datasets, transforms
  transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
  ds = datasets.MNIST(“./data”, train=True, download=True, transform=transform)
  train_ds, val_ds = random_split(ds, [int(0.9len(ds)), len(ds)-int(0.9len(ds))])
  train_loader = DataLoader(train_ds, batch_size=64, shuffle=True, num_workers=4, pin_memory=True)
  val_loader = DataLoader(val_ds, batch_size=64, num_workers=4, pin_memory=True)
• 模型、损失与优化器：
  import torch.nn as nn, torch.optim as optim
  class Net(nn.Module):
  def init(self): super().init(); self.conv1=nn.Conv2d(1,32,3,1); self.conv2=nn.Conv2d(32,64,3,1)
  self.pool=nn.MaxPool2d(2); self.fc1=nn.Linear(9216,128); self.fc2=nn.Linear(128,10)
  def forward(self,x):
  x=F.relu(self.conv1(x)); x=F.relu(self.conv2(x)); x=self.pool(x)
  x=torch.flatten(x,1); x=F.relu(self.fc1(x)); return self.fc2(x)
  device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)
  model = Net().to(device); criterion = nn.CrossEntropyLoss()
  optimizer = optim.Adam(model.parameters(), lr=1e-3)
• 训练循环与验证：
  for epoch in range(3):
  model.train(); tot=0
  for xb, yb in train_loader:
  xb, yb = xb.to(device), yb.to(device)
  optimizer.zero_grad(); out=model(xb); loss=criterion(out,yb); loss.backward(); optimizer.step(); tot+=loss.item()
  print(f"Epoch {epoch+1}: train loss {tot/len(train_loader):.4f}“)
  model.eval(); corr=0; with torch.no_grad():
  for xb, yb in val_loader:
  xb, yb = xb.to(device), yb.to(device)
  pred = model(xb).argmax(dim=1); corr += (pred==yb).sum().item()
  print(f"Val acc: {corr/len(val_ds):.4f}”)
• 保存与加载：
  torch.save(model.state_dict(), “model.pth”)

  推理时：model.load_state_dict(torch.load(“model.pth”, map_location=device)); model.eval()

• 可选：可视化与监控
  pip install tensorboard; from torch.utils.tensorboard import SummaryWriter; writer = SummaryWriter(“runs/exp1”); writer.add_scalar(“Loss/train”, loss.item(), epoch); writer.close(); 之后执行 tensorboard --logdir=runs。

四 部署与服务化

• 本地推理脚本要点：
  import torch, json
  from model import Net # 与训练一致的模型定义
  device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)
  model = Net().to(device); model.load_state_dict(torch.load(“model.pth”, map_location=device)); model.eval()
  def predict(x):
  with torch.no_grad(): inp=torch.tensor(x, dtype=torch.float32).unsqueeze(0).to(device)
  return model(inp).cpu().numpy().tolist()
• HTTP 服务示例（Flask）：
  from flask import Flask, request, jsonify
  app = Flask(name); predict = predict # 绑定上面的函数
  @app.route(“/predict”, methods=[“POST”])
  def s(): return jsonify(predict(request.json[“input”]))
  if name == “main”: app.run(host=“0.0.0.0”, port=5000)
• 运行与测试：
  nohup python app.py > serve.log 2>&1 &
  curl -X POST -H “Content-Type: application/json” -d ‘{“input”: [[0.0]*784]}’ http://localhost:5000/predict
• 生产建议：使用 Gunicorn/Uvicorn+FastAPI 承载高并发，容器化（Docker）封装运行环境，Nginx 反向代理与限流，加入日志、熔断与模型版本管理。

五 常见问题与排障清单

• 驱动与工具链：nvidia-smi 正常但 torch.cuda.is_available() 为 False，多为驱动与 PyTorch 的 CUDA 版本不匹配；优先使用 conda 的 cudatoolkit 或 pip 的官方 cu 版本，避免系统 CUDA 与 PyTorch 内置版本冲突。
• 环境变量：若需手动配置 CUDA，确保 ~/.bashrc 或系统 profile 中有
  export PATH=/usr/local/cuda/bin:$PATH
  export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
  并执行 source ~/.bashrc 生效。
• 权限与资源：多用户服务器上，确保用户对 /dev/nvidia* 有访问权限；使用 nvidia-smi 观察显存占用与进程，必要时设置 CUDA_VISIBLE_DEVICES 隔离 GPU。
• 稳定性：设置固定随机种子、开启 deterministic 选项（可能影响性能）、使用 pin_memory 与合适 num_workers 提升数据吞吐；OOM 时降低 batch_size 或使用梯度累积。