「AI&Tools」宿舂粮

自在即为仙收录于工程

2026-08-14 约 1223 字预计阅读 3 分钟

记录并整理人工智能领域相关工具的基础用法。

框架

PyTorch

多卡训练

在容器中使用多卡训练首先需要将显卡及驱动暴露到容器内部，有两种方案

暴露全部显卡，内部按需取用

仅暴露可用显卡，内部全部或按需取用

# 启动容器时指定显卡
docker run ... --gpus '"device=2,3,4"' ...  # 仅暴露指定卡，注意内部卡号会被从 `0` 开始重编

# 或容器内指定显卡
export CUDA_VISIBLE_DEVICES=2,3,4

在使用时也有几种方案

使用 DataParallel ，即在训练前给模型包上 nn.DataParallel

import torch
import torch.nn as nn

# 训练前
device_ids = [0, 1, 2]   # 因为 CUDA_VISIBLE_DEVICES=2,3,4
model = compile_detection_model(grid_conf, data_aug_conf)
model = nn.DataParallel(model, device_ids=device_ids).cuda()

# 训练时不再写单卡 device
preds = model(imgs.cuda(), rots.cuda(), trans.cuda(), intrins.cuda(), post_rots.cuda(), post_trans.cuda())

然后，运行如下训练命令

python main.py train_det mini/trainval --dataroot=/data/nuscenes --logdir=./runs_det --bsz=4

使用 DistributedDataParallel

model = model.to(local_rank)
model = DistributedDataParallel(model, device_ids=[local_rank])
DistributedSampler()

运行如下训练命令

export CUDA_VISIBLE_DEVICES=2,3,4
torchrun --nproc_per_node=3 main.py train_det mini/trainval --dataroot=/data/nuscenes --logdir=./runs_det --bsz=4
# 或单条命令方式（避免导入的 CUDA_VISIBLE_DEVICES 影响后续命令）
CUDA_VISIBLE_DEVICES=2,3,4 torchrun ...
# 或容器外直接执行
docker compose exec -e CUDA_VISIBLE_DEVICES=2,3,4 lss torchrun --standalone --nproc_per_node=3 ...

这两种方式的对比

维度	`DataParallel`	`DistributedDataParallel`
架构与进程	单进程、多线程	多进程
硬件支持	单机多卡	单机多卡或多机多卡
通信机制	每批次执行多次数据交换与广播	采用 Ring-AllReduce 模式，每批次仅在反向传播时同步一次梯度
显卡负载	不均衡主卡需要汇总梯度、分发数据，容易显存爆满	均衡各进程独立计算、梯度同步均摊到所有卡，无明显主卡瓶颈
训练效率	较低线程通信开销大和主卡瓶颈	很高接近线性的加速比
使用难度	简单（但过时）	稍复杂（标准方式）

TensorFlow

训练可视化与监控

Nvidia Driver

watch -n 1 nvidia-smi  # 每秒刷新 Nvidia 显卡状态

1. 查找占用该显卡的用户信息

# 确定占用该卡的进程 `PID`
nvidia-smi

# 显示该进程的进程树
pstree -aps <PID>
# -a: 显示启动每个进程对应的完整指令，包括启动进程的路径、参数等
# -c: 不使用精简法显示进程信息，即显示的进程中包含子进程和父进程
# -n: 根据进程 PID 号来排序输出（默认以程序名排序输出）
# -p: 显示进程的 PID
# -u: （当进程启动用户与其父进程不同时）显示进程对应的用户名称
# -s: 仅显示指定进程的所有祖先直到系统根进程

# 如果用户从容器启动，则查看该容器信息
# ... 若容器信息看不出来，可以进去看看代码等信息
docker ps -a | grep <ID前缀>

# 如果用户宿主机启动，则查询顶层 `bash` 或 `tmux` 等用户进程的 PID
ps -F -p <top PID>

TensorBoard

由 Google 开发的免费开源训练过程可视化工具。其完全本地运行，读取本地存储的数据能

实时绘制 Loss 和 Accuracy 折线图，直方图展示权重和梯度变化
展现神经网络的模型计算图和数据流向
即时渲染训练过程中的文本、图像、音频和 3D 点云等

启动

tensorboard --logdir=runs
# --port: 指定 Web 服务端口号（默认 6006）
# --host: 绑定指定 IP 地址，设为 0.0.0.0 可外网访问
# --bind_all: 一键开启全网段监听，便于局域网访问

# --reload_interval: 设置后台检测并同步新数据的秒数间隔（默认 5 秒）
# --samples_per_plugin: 限制前端加载的标量或图片的最大样本数，防止网页过载卡死
# --max_reload_threads: 设定多线程并行加载日志的线程数，加速超大日志读取

# --window_title: 自定义浏览器标签页显示的标题名称
# --purge_orphaned_data: 训练意外崩溃重启后，自动清理未保存的孤立数据（默认开启）

# 可搭配 SSH 端口转发实现本地安全访问服务器端口
ssh -p 22 -L 6006:localhost:6006 user@server_ip

数据记录

以 PyTorch 框架为例

from torch.utils.tensorboard import SummaryWriter

# 创建一个写入器，日志存放在 'runs/experiment_1'
writer = SummaryWriter('runs/experiment_1')

for epoch in range(100):
    # 模拟训练过程
    train_loss = 1.0 / (epoch + 1)

    # 核心：写入标量数据
    writer.add_scalar('Loss/train', train_loss, epoch)

# 记得关闭写入器
writer.close()

Weights & Biases (W&B)

目前业界最流行的商业化试验追踪平台。

MLflow

由 Databricks 主导的开源平台，注重机器学习全流程的规范化。

目录