Compare commits

...
3 Commits
Author SHA1 Message Date
Shaun ZhouandClaude Opus 5 938d62ea53 docs: SimCLR 域内自监督的实测对比与多后端说明
- 新增 §10.5:SimCLR vs ImageNet 预训练在 14,579 张上的同设置对比。
  只训 1 个 epoch 的域内特征在 k=4 上即追平(轮廓 0.334 vs 0.325),
  在 k=2 上大幅胜出(轮廓 0.469、稳定性 0.991)——目视确认该二分正是
  "有明显缺陷 / 弱纹理近空白",可直接作为流水线的第一级过滤
- 复杂度表补 MPS 实测:特征提取 84s(CPU) -> 51s(MPS),SimCLR 约 65s/epoch
- README 增加"计算后端"小节(三后端的 AMP 默认值、显存统计与限制)
- 流程图与对照表同步去掉已移除的 GPU 守卫

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-02 21:43:35 +08:00
Shaun ZhouandClaude Opus 5 b6f83fcad0 fix(simclr): 周期性落盘,中途中断不再丢失训练成果
此前 fit() 只在全部 epoch 跑完后才保存骨架,一次跑到 epoch 12 被停掉后
13 分钟的训练全部作废。

- 每 save_interval(默认 5)个 epoch 保存一次骨架
- 捕获 KeyboardInterrupt,中断时先落盘再返回
- 聚类脚本改用 resolve_device,自监督分支不再有 GPU 守卫

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-02 21:43:35 +08:00
Shaun ZhouandClaude Opus 5 141cb907f2 feat: 支持 cuda / mps / cpu 三后端,Mac 上可用 MPS 训练
移除"训练必须有 Nvidia GPU"的硬性守卫,改为按后端能力分派。

- utils/device: resolve_device("auto") 按 cuda > mps > cpu 挑最快的可用后端;
  显式指定不可用的后端时直接报错,不会悄悄退回 CPU——训练慢几十倍却不吭声是最坏的情况
- AMP 随后端取默认值:CUDA 开、MPS 关。MPS 上 autocast 数值正常但实测更慢
  (tiny_conv @128²:fp32 76 ms/步、fp16 90 ms、bf16 108 ms),需要时 amp: true 可强开
- PeakMemory 统一显存峰值统计:CUDA 读内置计数器,MPS 无峰值 API 改用逐 step 采样
  (是下界,日志与产出里已标注)
- profile_inference 支持 MPS,并提示 torch.profiler 没有 MPS 的 ProfilerActivity,
  算子明细只反映 CPU 侧下发耗时
- train.py / eval.py 新增 --device,配置默认改为 auto
- check_env 改为列出全部可用后端与 auto 的选择结果

实测:tiny_conv 训练在 MPS 上比 CPU 快约 32 倍。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-02 21:43:35 +08:00
22 changed files with 630 additions and 262 deletions
+34 -10
View File
@@ -25,13 +25,30 @@ convolution 的结构(不含 attention / transformer),并支持导出 ONNX
```bash
uv sync # 安装依赖(含 swanlab)
uv sync --group export # 需要导出 ONNX 时再装
uv run python scripts/check_env.py # 自检(Python / PyTorch / Nvidia GPU)
uv run python scripts/check_env.py # 自检(Python / PyTorch / 可用后端)
```
> **训练机必须有 Nvidia GPU。** 所有训练入口先调用 `utils.device.require_cuda()`,
> 检测不到 CUDA 构建的 PyTorch 或可用 GPU 时直接报错退出。评估与 ONNX 导出允许 CPU。
> Linux 上 PyPI 的 torch 即 CUDA 构建;如需指定 CUDA 版本:
> `uv pip install --index-url https://download.pytorch.org/whl/cu124 torch torchvision`
### 计算后端
支持 **CUDA / MPS(Apple Silicon)/ CPU**。配置里的 `device: auto` 按 `cuda > mps > cpu`
自动挑最快的可用后端,也可以显式写 `cuda` / `mps` / `cpu`,或用命令行 `--device` 覆盖:
```bash
uv run python train.py --config configs/train/baseline.yaml # auto
uv run python train.py --config configs/train/baseline.yaml --device mps # 强制 Mac GPU
uv run python train.py --config configs/train/baseline.yaml --device cpu
```
**显式指定不可用的后端会直接报错退出**,不会悄悄退回 CPU——训练慢几十倍却不吭声是最坏的情况。
| 后端 | AMP 默认 | 显存统计 | 备注 |
| --- | --- | --- | --- |
| `cuda` | **开** | 精确峰值 | Linux 上 PyPI 的 torch 即 CUDA 构建;指定版本用 `uv pip install --index-url https://download.pytorch.org/whl/cu124 torch torchvision` |
| `mps` | **关** | 采样下界 | Apple Silicon。autocast 可用且数值正常,但实测比 fp32 慢(tiny_conv @128²:fp32 76 ms/步,fp16 90 ms),故默认不开;`amp: true` 可强制打开 |
| `cpu` | 不支持 | 无 | 实测比 MPS 慢约 30 倍,仅适合调试与精度回归 |
`amp` 字段接受 `auto`(随后端)/ `true` / `false`。MPS 上 `torch.profiler` 抓不到 Metal kernel
(没有对应的 `ProfilerActivity`),算子明细只反映 CPU 侧下发耗时。
## 数据准备
@@ -148,7 +165,7 @@ uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml
# 先扫一遍簇数再决定分几类
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml --sweep-k 2-8
# 换成在本地数据上自监督训练的特征(SimCLR,属于训练,需要 Nvidia GPU)
# 换成在本地数据上自监督训练的特征(SimCLR,Mac 上自动走 MPS)
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml --set method=simclr
```
@@ -177,8 +194,8 @@ uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml --
**本批薄膜图上自监督训练**一个项目自己的纯卷积骨架,域适配更好。同一张图的两个随机增强
视图互为正样本,batch 内其余 2B−2 个视图为负样本,用 NT-Xent 损失训练。
**这属于训练,因此必须有 Nvidia GPU**(走 `utils.device.require_cuda()` 守卫,
在无卡机器上会直接报错退出):
训练走 `device: auto` 选中的后端——**Mac 上直接用 MPS 即可**,实测 14,579 张、
batch 256 时约 65 秒/epoch:
```bash
# 1. 自监督训练 + 聚类一步跑完(首次运行)
@@ -191,7 +208,13 @@ uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml --set method=simclr simclr.epochs=200 simclr.batch_size=512 simclr.temperature=0.1
```
训练完权重存到 `<work_dir>/simclr_backbone.pth`(**只含骨架,不含投影头**——投影头被训练成
**实测效果**(14,579 张,与 pretrained 用完全相同的聚类设置):只训 **1 个 epoch** 的域内特征
在 k=4 上就追平 ImageNet(轮廓 0.334 vs 0.325),在 **k=2 上大幅胜出**(轮廓 0.469、
稳定性 0.991)——目视确认这个二分正是"有明显缺陷 / 弱纹理近空白",可直接当作流水线的
第一级过滤。详见 [文档 §10.5](docs/unsupervised-clustering.md#105-simclr-域内自监督-vs-imagenet-预训练)。
训练每 `save_interval` 个 epoch 落一次盘,中途 Ctrl-C 也能拿到当前权重继续用。
权重存到 `<work_dir>/simclr_backbone.pth`(**只含骨架,不含投影头**——投影头被训练成
对增强不变,会丢掉下游有用的信息,所以聚类和保存都只用骨架特征)。相关配置:
```yaml
@@ -204,7 +227,8 @@ simclr:
epochs: 100
batch_size: 256
temperature: 0.2 # 越小越关注难负样本
amp: true
amp: auto # CUDA 上开、MPS 上关
save_interval: 5 # 每 5 个 epoch 落一次盘
checkpoint: null # 填已有权重可跳过重复训练
projection: {hidden_channels: 256, out_channels: 128}
optimizer: {name: adamw, lr: 0.001, weight_decay: 0.0001}
+5 -4
View File
@@ -22,12 +22,12 @@
work_dir: runs/cluster
seed: 42
device: cuda # method=pretrained 时可以 --set device=cpu
device: auto # auto 按 cuda > mps > cpu 挑最快的可用后端;也可写 cuda / mps / cpu
num_clusters: 6 # 想分几类改这里(或 --num-clusters 6)
num_clusters: 4 # 想分几类改这里(或 --num-clusters 6)
# pretrained: ImageNet 预训练卷积骨架直接抽特征,不训练,CPU 也能跑,几分钟出结果。
# simclr: 在本目录数据上自监督训练项目自己的卷积骨架,域适配更好,需要 Nvidia GPU。
# simclr: 在本目录数据上自监督训练项目自己的卷积骨架,域适配更好;Mac 上走 MPS 即可。
method: pretrained
# method=pretrained 用的骨架(method=simclr 的骨架配在 simclr.backbone,两者字段完全不同,
@@ -97,7 +97,8 @@ simclr:
epochs: 100
batch_size: 256
temperature: 0.2 # 越小越关注难负样本
amp: true
amp: auto # CUDA 上开、MPS 上关
save_interval: 5 # 每 5 个 epoch 落一次盘,中途中断也能拿到权重
log_interval: 20
checkpoint: null # 填已有的 simclr_backbone.pth 可跳过重复训练
projection: {hidden_channels: 256, out_channels: 128}
+2 -2
View File
@@ -1,8 +1,8 @@
_base_:
- ../train/baseline.yaml
# 评估允许在 CPU 上跑(部署前的精度回归),训练则必须用 GPU。
device: cuda
# auto 按 cuda > mps > cpu 自动选;部署前做精度回归可以显式 device=cpu。
device: auto
checkpoint: runs/baseline/best.pth
eval:
+2 -2
View File
@@ -5,11 +5,11 @@ _base_:
work_dir: runs/baseline
seed: 42
deterministic: false
device: cuda # 训练只允许 Nvidia GPU,见 utils/device.py
device: auto # auto 按 cuda > mps > cpu 挑最快的可用后端;也可写 cuda / mps / cpu
train:
epochs: 100
amp: true
amp: auto # auto: CUDA 上开、MPS 上关(实测 MPS autocast 比 fp32 慢);也可写 true / false
clip_grad: 5.0
log_interval: 20
monitor: f1 # 用于挑选 best.pth 的验证指标
+3 -3
View File
@@ -5,11 +5,11 @@ from __future__ import annotations
from collections.abc import Mapping
from typing import Any
import torch
from torch.utils.data import DataLoader, Dataset
from datasets.film_defect import collate_fn
from datasets.registry import DATASETS
from utils.device import pin_memory_supported, resolve_device
def build_dataset(cfg: Mapping[str, Any], **kwargs: Any) -> Dataset:
@@ -24,9 +24,9 @@ def build_dataloader(
pin_memory: bool | None = None,
drop_last: bool = False,
) -> DataLoader:
"""`pin_memory` 缺省时只在有 CUDA 设备时开启(CPU/MPS 上开启无意义且会告警)。"""
"""`pin_memory` 缺省时只在 CUDA 上开启(CPU/MPS 上无收益且会告警)。"""
if pin_memory is None:
pin_memory = torch.cuda.is_available()
pin_memory = pin_memory_supported(resolve_device("auto"))
return DataLoader(
dataset,
batch_size=batch_size,
+56 -71
View File
@@ -20,7 +20,7 @@
7. [阶段五:无标签评估](#7-阶段五无标签评估)
8. [阶段六:SimCLR 自监督(可选)](#8-阶段六simclr-自监督可选)
9. [复杂度与实测耗时](#9-复杂度与实测耗时)
10. [消融实验数据](#10-消融实验数据)
10. [消融实验数据](#10-消融实验数据)(含 [SimCLR vs ImageNet 对比](#105-simclr-域内自监督-vs-imagenet-预训练))
11. [已知局限与改进方向](#11-已知局限与改进方向)
12. [公式与代码对照表](#12-公式与代码对照表)
@@ -130,7 +130,7 @@ flowchart TB
summ["summary.json<br/>簇大小 / 轮廓 / 稳定性"]:::out
end
simclr["可选:SimCLR 自监督<br/>method = simclr<br/>需 Nvidia GPU"]:::feat
simclr["可选:SimCLR 自监督<br/>method = simclr<br/>cuda / mps 均可"]:::feat
scan --> PRE
PRE -->|"x̃ᵢ"| FEAT
@@ -330,8 +330,7 @@ xychart
| 32 | 0.922 | 0.050 | 59% |
| 64 | 0.922 | 0.047 | 82% |
不白化时结果对 $q$ 几乎不敏感(都是 0.92),白化则全线崩溃。在真实数据上,
早期开白化的版本最大簇吃掉 **96%** 的样本、稳定性仅 0.01;关掉后最大簇降到 38%。
不白化时结果对 $q$ 几乎不敏感(都是 0.92),白化则全线崩溃。在真实数据上,早期开白化的版本最大簇吃掉 **96%** 的样本、稳定性仅 0.01;关掉后最大簇降到 38%。
> 这条结论被回归测试 `test_whitening_amplifies_noise_directions` 锁死。
@@ -405,16 +404,13 @@ $$
\mu_j = \frac{1}{|C_j|}\sum_{i \in C_j} y_i .
$$
由于 $J$ 在两步中都不增、且 $J \ge 0$ 有下界,同时 $N$ 个点划分成 $k$ 簇的方式**有限**,
算法必在有限步内停止。收敛判据用相对变化:$|J_{\text{prev}} - J| \le \text{tol}\cdot J_{\text{prev}}$,
`tol = 1e-5`,上限 100 次迭代。
由于 $J$ 在两步中都不增、且 $J \ge 0$ 有下界,同时 $N$ 个点划分成 $k$ 簇的方式**有限**,算法必在有限步内停止。收敛判据用相对变化:$|J_{\text{prev}} - J| \le \text{tol}\cdot J_{\text{prev}}$,`tol = 1e-5`,上限 100 次迭代。
> **注意**:收敛保证的是局部最优,不是全局最优——这正是需要 `n_init` 的原因。
### 6.3 k-means++ 初始化
均匀随机取初值时,多个初始质心可能落在同一团里,导致明显的坏解。k-means++ 按到已选
质心的平方距离采样:
均匀随机取初值时,多个初始质心可能落在同一团里,导致明显的坏解。k-means++ 按到已选质心的平方距离采样:
$$
\Pr(\text{选中 } y) = \frac{D(y)^2}{\sum_{y'} D(y')^2},
@@ -431,28 +427,22 @@ $$
### 6.4 空簇处理
更新步中若某簇没有任何成员,$\mu_j$ 无定义。实现的处理是:**把当前离自身质心最远的点
提为该空簇的新质心**。这保证输出恰好 $k$ 个非空簇(否则实际簇数会悄悄缩水),且由于选的是
$J$ 贡献最大的点,这一步通常还会降低 $J$。
更新步中若某簇没有任何成员,$\mu_j$ 无定义。实现的处理是:**把当前离自身质心最远的点提为该空簇的新质心**。这保证输出恰好 $k$ 个非空簇(否则实际簇数会悄悄缩水),且由于选的是 $J$ 贡献最大的点,这一步通常还会降低 $J$。
### 6.5 多次重启
跑 `n_init = 10` 次完整的(初始化 + Lloyd),取 $J$ 最小的一次。设单次落入某个坏局部最优的
概率为 $p$,则 10 次全坏的概率降到 $p^{10}$。
跑 `n_init = 10` 次完整的(初始化 + Lloyd),取 $J$ 最小的一次。设单次落入某个坏局部最优的概率为 $p$,则 10 次全坏的概率降到 $p^{10}$。
### 6.6 子集拟合 + 全量指派
可选地只用随机子集 $\mathcal{S} \subset \{1..N\}$(`fit_max_samples`)拟合质心,再把全部
$N$ 个样本按最近质心指派:
可选地只用随机子集 $\mathcal{S} \subset \{1..N\}$(`fit_max_samples`)拟合质心,再把全部 $N$ 个样本按最近质心指派:
$$
\mu^* = \arg\min_\mu J_{\mathcal{S}}(\mu), \qquad
c(i) = \arg\min_j \lVert y_i - \mu^*_j \rVert, \ \ i = 1,\dots,N .
$$
合理性:质心是簇内均值的估计,其标准误随 $1/\sqrt{|\mathcal{S}|}$ 下降,几千个样本已足够。
**子集必须是纯随机抽取**——按目录或任何人工整理的历史标签做分层,等于把先验偷偷注入
无监督流程。默认 `fit_max_samples: 0`(用全部样本),因为 14.5k × 32 维的 K-Means 本来就只要几秒。
合理性:质心是簇内均值的估计,其标准误随 $1/\sqrt{|\mathcal{S}|}$ 下降,几千个样本已足够。**子集必须是纯随机抽取**——按目录或任何人工整理的历史标签做分层,等于把先验偷偷注入无监督流程。默认 `fit_max_samples: 0`(用全部样本),因为 14.5k × 32 维的 K-Means 本来就只要几秒。
---
@@ -476,8 +466,7 @@ s(i) = \frac{b(i) - a(i)}{\max\{a(i),\, b(i)\}} \in [-1, 1],
\bar{s} = \frac{1}{N}\sum_i s(i).
$$
$s(i) \to 1$ 表示样本牢牢待在自己簇里,$s(i) < 0$ 表示它离别的簇更近(八成分错了)。
计算是 $O(N^2 d)$,故超过 `max_samples = 2000` 时随机抽样估计。
$s(i) \to 1$ 表示样本牢牢待在自己簇里,$s(i) < 0$ 表示它离别的簇更近(八成分错了)。计算是 $O(N^2 d)$,故超过 `max_samples = 2000` 时随机抽样估计。
### 7.2 换种子稳定性(划分是否可复现)
@@ -494,12 +483,9 @@ I(\mathcal{A};\mathcal{B}) = \sum_{a}\sum_{b} p(a,b)\log\frac{p(a,b)}{p(a)p(b)},
H(\mathcal{A}) = -\sum_a p(a)\log p(a),
$$
其中 $p(a,b) = |A_a \cap B_b| / N$。NMI 对簇编号的置换不变(换标签名不影响),
取值 $[0,1]$,1 表示两次划分完全一致。
其中 $p(a,b) = |A_a \cap B_b| / N$。NMI 对簇编号的置换不变(换标签名不影响),取值 $[0,1]$,1 表示两次划分完全一致。
**为什么这个指标关键**:数据里真有 $k$ 团结构时,不同初始化会收敛到同一套划分,NMI ≈ 1;
若是在硬切一团连续的数据,边界位置随初值乱跳,NMI 会很低。它比轮廓系数更能识破
"强行分 $k$ 类"的情形。实测中它一眼揪出了白化的问题(稳定性 0.01)。
**为什么这个指标关键**:数据里真有 $k$ 团结构时,不同初始化会收敛到同一套划分,NMI ≈ 1;若是在硬切一团连续的数据,边界位置随初值乱跳,NMI 会很低。它比轮廓系数更能识破"强行分 $k$ 类"的情形。实测中它一眼揪出了白化的问题(稳定性 0.01)。
> **不使用外部标签。** `data/film` 下的目录名(`grease`、`wrinkle`、`反面纵黑点`……)
> 是人工整理的历史产物,既不完整也未必正确,因此既不参与聚类、也不用于评估。
@@ -508,9 +494,7 @@ $$
## 8. 阶段六:SimCLR 自监督(可选)
`method: pretrained` 用的是 ImageNet 权重——域外的。`method: simclr` 则在**这批薄膜图
本身**上学特征,域适配更好。它是训练,因此按项目约定必须有 Nvidia GPU
(`utils.device.require_cuda()` 守卫)。
`method: pretrained` 用的是 ImageNet 权重——域外的。`method: simclr` 则在**这批薄膜图本身**上学特征,域适配更好。它是训练,设备由 `device: auto` 按 `cuda > mps > cpu` 自动挑选:CUDA 与 Apple Silicon 的 MPS 都能跑(实测 14,579 张、batch 256 时 MPS 约 65 秒/epoch),CPU 慢约 30 倍不建议。
```mermaid
flowchart TB
@@ -572,12 +556,9 @@ $$
x_i^{(1)} = t(I_i), \qquad x_i^{(2)} = t'(I_i).
$$
$\mathcal{T}$ 包含 `RandomResizedCrop(scale=0.5~1.0)`、水平/垂直翻转、高斯模糊。
**监督信号完全来自"这两个视图来自同一张图"这一事实**,不需要任何人工标签。
$\mathcal{T}$ 包含 `RandomResizedCrop(scale=0.5~1.0)`、水平/垂直翻转、高斯模糊。**监督信号完全来自"这两个视图来自同一张图"这一事实**,不需要任何人工标签。
> `scale` 下限不宜低于 0.4:裁得太狠容易两个视图都只剩背景,"背景 vs 背景"被当成正样本对,
> 学出的特征会退化。另外主 transforms 里的 `PerImageStandardize` 在增强**之后**执行,
> 因此 `ColorJitter` 的亮度扰动最终会被抵消,真正起作用的是裁剪/翻转/模糊。
> `scale` 下限不宜低于 0.4:裁得太狠容易两个视图都只剩背景,"背景 vs 背景"被当成正样本对,学出的特征会退化。另外主 transforms 里的 `PerImageStandardize` 在增强**之后**执行,因此 `ColorJitter` 的亮度扰动最终会被抵消,真正起作用的是裁剪/翻转/模糊。
### 8.2 NT-Xent 损失
@@ -596,16 +577,11 @@ $$
\log \frac{\exp\!\big(S_{k,\pi(k)}\big)}{\sum_{m \neq k}\exp\!\big(S_{k,m}\big)}\ }
$$
分子只有 1 个正样本对,分母是全部 $2B-2$ 个负样本加正样本对。形式上就是把
"在 $2B-1$ 个候选中认出配对视图"当成一个 $(2B-1)$ 类的分类问题——实现里直接用
`F.cross_entropy(S, targets)`,`targets[k] = π(k)`。
分子只有 1 个正样本对,分母是全部 $2B-2$ 个负样本加正样本对。形式上就是把"在 $2B-1$ 个候选中认出配对视图"当成一个 $(2B-1)$ 类的分类问题——实现里直接用 `F.cross_entropy(S, targets)`,`targets[k] = π(k)`。
**温度 $\tau$** 控制分布的尖锐度:$\tau$ 小 ⇒ softmax 更尖 ⇒ 梯度集中在最难的负样本上;
$\tau$ 大 ⇒ 所有负样本近乎等权。默认 $\tau = 0.2$。
**与互信息的关系**:NT-Xent 是 InfoNCE 的一个实例,最小化 $\mathcal{L}$ 等价于最大化两个
视图表征之间互信息的一个下界:
**温度 $\tau$** 控制分布的尖锐度:$\tau$ 小 ⇒ softmax 更尖 ⇒ 梯度集中在最难的负样本上;$\tau$ 大 ⇒ 所有负样本近乎等权。默认 $\tau = 0.2$。
**与互信息的关系**:NT-Xent 是 InfoNCE 的一个实例,最小化 $\mathcal{L}$ 等价于最大化两个视图表征之间互信息的一个下界:
$$
I\big(h^{(1)}; h^{(2)}\big) \;\ge\; \log(2B-1) - \mathcal{L}.
$$
@@ -614,16 +590,12 @@ $$
### 8.3 为什么聚类用骨架特征而不是投影头输出
$g_\varphi$ 被显式训练成对增强 $\mathcal{T}$ **不变**。凡是被增强破坏掉的信息(模糊掉的
高频细节、翻转掉的方向性),在 $z$ 里都会被压掉。而这些信息对下游任务往往有用。
$h = f_\theta(x)$ 位于投影头之前,保留得更完整——这是 SimCLR 原论文的既有结论
(线性评估上 $h$ 比 $z$ 高十几个点)。因此:**训练用 $z$ 算损失,聚类和保存权重都只用 $h$。**
$g_\varphi$ 被显式训练成对增强 $\mathcal{T}$ **不变**。凡是被增强破坏掉的信息(模糊掉的高频细节、翻转掉的方向性),在 $z$ 里都会被压掉。而这些信息对下游任务往往有用。
$h = f_\theta(x)$ 位于投影头之前,保留得更完整——这是 SimCLR 原论文的既有结论(线性评估上 $h$ 比 $z$ 高十几个点)。因此:**训练用 $z$ 算损失,聚类和保存权重都只用 $h$。**
### 8.4 与检测任务的衔接
`simclr_backbone.pth` 里的骨架就是项目检测模型用的同一个卷积骨架,可以直接拿去初始化
`DefectNet` 的 backbone——在标注量有限时,用 14.5k 张无标注图预训练过的权重通常比
ImageNet 权重更贴合薄膜的成像特性。
`simclr_backbone.pth` 里的骨架就是项目检测模型用的同一个卷积骨架,可以直接拿去初始化 `DefectNet` 的 backbone——在标注量有限时,用 14.5k 张无标注图预训练过的权重通常比 ImageNet 权重更贴合薄膜的成像特性。
---
@@ -641,7 +613,7 @@ ImageNet 权重更贴合薄膜的成像特性。
| 产出(软链 + 拼图) | $O(N + k\cdot\text{preview})$ | ~2 s |
| **合计** | | **~100 s** |
GPU 上特征提取会快一个数量级;瓶颈会转移到图像解码(可增大 `loader.num_workers`)。
换到 Apple Silicon 的 MPS(`device: auto` 自动选中)后,同样 14,579 张的特征提取从 84 秒降到 51 秒——瓶颈已转移到图像解码,继续加速要靠增大 `loader.num_workers`。SimCLR 自监督训练在 MPS 上约 65 秒/epoch(batch 256,双视图,显存峰值约 137 MB)。
---
@@ -688,10 +660,32 @@ GPU 上特征提取会快一个数量级;瓶颈会转移到图像解码(可
| 7 | 0.230 | 0.692 | 0.28 | [236, 345, 87, 509, 559, 150, 114] |
| 8 | 0.272 | 0.758 | 0.34 | [676, 215, 58, 27, 176, 99, 430, 319] |
$k=4$ 的稳定性在 $k \in [3,7]$ 中最高,轮廓也接近峰值——默认值 4 有数据支撑。
命令行随时可扫:`--sweep-k 2-8`。
$k=4$ 的稳定性在 $k \in [3,7]$ 中最高,轮廓也接近峰值——默认值 4 有数据支撑。命令行随时可扫:`--sweep-k 2-8`。
### 10.5 全量最终结果
### 10.5 SimCLR 域内自监督 vs ImageNet 预训练
在全部 14,579 张上、用**完全相同**的聚类设置(PCA 32 维、不白化、`n_init=10`)对比两种特征。
SimCLR 一侧用的是在本批数据上**只训了 1 个 epoch**(Apple M 系列 MPS,65 秒)的 tiny_conv 骨架:
| 特征 | 维度 | $k=4$ 轮廓 | $k=4$ 稳定性 | $k=2$ 轮廓 | $k=2$ 稳定性 |
| --- | --- | --- | --- | --- | --- |
| ImageNet 预训练 layer2+3 | 384 | 0.325 | **0.977** | 0.354 | 0.887 |
| SimCLR 域内(1 epoch) | 256 | **0.334** | 0.829 | **0.469** | **0.991** |
两点值得注意:
1. **只训 1 个 epoch 就在 $k=4$ 上追平并略微超过 ImageNet 特征**(0.334 vs 0.325),而它的维度还少三分之一。域内自监督的性价比很高。
2. **$k=2$ 上 SimCLR 明显更强**(轮廓 0.469、稳定性 0.991)。目视确认这个二分就是"有明显缺陷 / 弱纹理近空白":8,750 张落在空白侧,5,829 张落在缺陷侧(褶皱、V 形折痕、条纹、异物、点串)。这正是 [§11](#11-已知局限与改进方向) 提到的"先筛掉无缺陷样本再细分"那一步——SimCLR 特征直接把它做出来了,可以作为流水线的第一级过滤。
$k=4$ 时 SimCLR 的四个簇为:暗底细长条纹/点串/异物(2,423)、空白暗底两类(5,801 + 2,927)、亮底团块/污渍(3,428)。相比 ImageNet 特征的划分,两个缺陷簇的代表样本明显更同质。
$k=4$ 稳定性略低于 ImageNet(0.83 vs 0.98),符合"数据的主结构是二分、强行切四份时边界不唯一"的解释。
> **这是下界,不是 SimCLR 的上限。** 原计划跑 40 epoch,中途在 epoch 12(loss 4.55 → 2.66)停掉,
> 而当时 `fit()` 只在全部 epoch 结束后落盘,训练成果没保住。该缺陷已修复(`save_interval`
> 每 5 个 epoch 存一次,并捕获 KeyboardInterrupt),上表用的是此前完整跑完的 1-epoch 权重。
### 10.6 全量最终结果
| 指标 | 值 |
| --- | --- |
@@ -715,25 +709,15 @@ $k=4$ 的稳定性在 $k \in [3,7]$ 中最高,轮廓也接近峰值——默
## 11. 已知局限与改进方向
1. **两个簇是"近乎空白"的。** 这批裁片里确实有很大比例在 128×128 尺度下看不出明显缺陷,
$k=4$ 就会把其中两个名额花在它们的细微差异上。这是数据本身的性质,不是算法缺陷。
对策:调大 $k$ 让缺陷簇进一步细分;或先用一个简单的对比度/能量阈值把"无明显缺陷"
的样本筛掉,再对剩下的做聚类。
1. **两个簇是"近乎空白"的。** 这批裁片里确实有很大比例在 128×128 尺度下看不出明显缺陷,$k=4$ 就会把其中两个名额花在它们的细微差异上。这是数据本身的性质,不是算法缺陷。对策:调大 $k$ 让缺陷簇进一步细分;或先用一个简单的对比度/能量阈值把"无明显缺陷"的样本筛掉,再对剩下的做聚类。
2. **ImageNet 特征终究是域外的。** `method: simclr` 在本批数据上自监督预训练,理论上更贴合,
但需要 GPU 且要调 epoch / batch size。建议在 GPU 机器上跑一次并与 pretrained 对比
轮廓系数与稳定性。
2. **ImageNet 特征终究是域外的。** 已实测:只训 1 个 epoch 的域内 SimCLR 就在 $k=4$ 上追平(0.334 vs 0.325),在 $k=2$ 上大幅胜出(0.469 / 稳定性 0.991),见 [§10.5](#105-simclr-域内自监督-vs-imagenet-预训练)。继续训到几十个 epoch 值得一试。
3. **拉伸带来的各向异性形变**在极端长条(40:1)上很严重。若发现某个簇专门收集长条,
可换 `CenterCrop`(原分辨率取中心 patch),或改成多 patch 聚合(对每张图取若干 patch
分别编码后取平均/最大)。后者对"缺陷是局部的"这一先验更契合,目前尚未实现。
3. **拉伸带来的各向异性形变**在极端长条(40:1)上很严重。若发现某个簇专门收集长条,可换 `CenterCrop`(原分辨率取中心 patch),或改成多 patch 聚合(对每张图取若干 patch 分别编码后取平均/最大)。后者对"缺陷是局部的"这一先验更契合,目前尚未实现。
4. **K-Means 假设簇是各向同性的球。** 若缺陷类别在特征空间里呈细长流形,K-Means 会切错。
备选:谱聚类(构图 + 拉普拉斯特征向量)或 HDBSCAN(不需要预设 $k$,能标记噪声点)。
代价是复杂度更高、超参更难调。
4. **K-Means 假设簇是各向同性的球。** 若缺陷类别在特征空间里呈细长流形,K-Means 会切错。备选:谱聚类(构图 + 拉普拉斯特征向量)或 HDBSCAN(不需要预设 $k$,能标记噪声点)。代价是复杂度更高、超参更难调。
5. **轮廓系数 0.33 属于"弱但存在"的结构。** 不要指望聚类结果能直接当标签用;
它的定位是**把人工核对的工作量从 14,579 张降到 4 个簇的抽查**。
5. **轮廓系数 0.33 属于"弱但存在"的结构。** 不要指望聚类结果能直接当标签用;它的定位是**把人工核对的工作量从 14,579 张降到 4 个簇的抽查**。
---
@@ -757,11 +741,12 @@ $k=4$ 的稳定性在 $k \in [3,7]$ 中最高,轮廓也接近峰值——默
| NT-Xent 损失 | `engine/simclr.py::nt_xent_loss` |
| SimCLR 训练循环 | `engine/simclr.py::SimCLRTrainer` |
| 投影头 $g_\varphi$ | `models/heads/projection.py::ProjectionHead` |
| GPU 守卫 | `utils/device.py::require_cuda` |
| 设备解析(cuda/mps/cpu,auto 优先级) | `utils/device.py::resolve_device` |
| AMP 随后端的默认值 | `utils/device.py::resolve_amp` / `default_amp` |
| 跨后端显存峰值 | `utils/device.py::PeakMemory` |
| CLI 与产出 | `scripts/cluster_unlabeled.py` |
**相关测试**:`tests/test_cluster.py`(K-Means 正确性、白化回归、稳定性、NMI)、
`tests/test_unlabeled_dataset.py`(扫描、混合尺寸/通道、逐图标准化的不变性)。
**相关测试**:`tests/test_cluster.py`(K-Means 正确性、白化回归、稳定性、NMI)、`tests/test_unlabeled_dataset.py`(扫描、混合尺寸/通道、逐图标准化的不变性)。
---
+44 -13
View File
@@ -1,7 +1,7 @@
"""SimCLR 自监督预训练:在无标注的薄膜图上学一套特征,供聚类使用。
骨架仍然是项目里的纯卷积 backbone(`tiny_conv` 等),预训练完可以直接拿去
初始化检测模型。属于"训练",因此调用方必须先过 `utils.device.require_cuda()`。
初始化检测模型。支持 cuda / mps / cpu,用 `device: auto` 自动挑最快的后端。
"""
from __future__ import annotations
@@ -19,6 +19,7 @@ from torch.utils.data import DataLoader
from engine.metrics import AverageMeter
from engine.optim import build_optimizer, build_scheduler
from models.heads.projection import ProjectionHead
from utils.device import PeakMemory, resolve_amp, synchronize
from utils.logging import get_logger
logger = get_logger("engine.simclr")
@@ -44,7 +45,7 @@ class SimCLRTrainer:
"""Args:
backbone: 输出特征图的卷积骨架(需带 `out_channels`)。
dataloader: 每个样本产出 2 个视图的 DataLoader(`UnlabeledImageDataset(views=2)`)。
cfg: `epochs / temperature / optimizer / scheduler / amp / log_interval` 等。
cfg: `epochs / temperature / optimizer / scheduler / amp / log_interval / save_interval` 等。
"""
def __init__(
@@ -65,6 +66,8 @@ class SimCLRTrainer:
self.epochs = int(self.cfg.get("epochs", 50))
self.temperature = float(self.cfg.get("temperature", 0.2))
self.log_interval = int(self.cfg.get("log_interval", 20))
# 每隔若干 epoch 落一次盘:自监督动辄跑几十分钟,中途被打断不该把训练成果全丢掉
self.save_interval = int(self.cfg.get("save_interval", 5))
self.projector = ProjectionHead(
backbone.out_channels, **dict(self.cfg.get("projection", {}))
@@ -74,8 +77,9 @@ class SimCLRTrainer:
self.model, self.cfg.get("optimizer", {"name": "adamw", "lr": 1e-3, "weight_decay": 1e-4})
)
self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs)
self.amp = bool(self.cfg.get("amp", True)) and device.type == "cuda"
self.scaler = torch.amp.GradScaler("cuda", enabled=self.amp)
self.amp = resolve_amp(self.cfg.get("amp", "auto"), device)
self.scaler = torch.amp.GradScaler(device.type, enabled=self.amp)
self.peak_memory = PeakMemory(device)
def train_one_epoch(self, epoch: int) -> float:
self.model.train()
@@ -85,7 +89,7 @@ class SimCLRTrainer:
for step, batch in enumerate(self.dataloader):
view1, view2 = (view.to(self.device, non_blocking=True) for view in batch["views"][:2])
with torch.amp.autocast("cuda", enabled=self.amp):
with torch.amp.autocast(self.device.type, enabled=self.amp):
z1 = self.projector(self.backbone(view1))
z2 = self.projector(self.backbone(view2))
loss = nt_xent_loss(z1, z2, self.temperature)
@@ -95,6 +99,7 @@ class SimCLRTrainer:
self.scaler.step(self.optimizer)
self.scaler.update()
self.peak_memory.update()
meter.update(float(loss.detach()), view1.shape[0])
if step % self.log_interval == 0:
logger.info(
@@ -106,20 +111,46 @@ class SimCLRTrainer:
self.optimizer.param_groups[0]["lr"],
)
logger.info("simclr epoch %d 完成 loss=%.4f 耗时 %.1fs", epoch, meter.avg, time.time() - start)
synchronize(self.device)
peak = self.peak_memory.peak_mb
logger.info(
"simclr epoch %d 完成 loss=%.4f 耗时 %.1fs 显存峰值 %.0f MB%s",
epoch,
meter.avg,
time.time() - start,
peak,
"" if self.peak_memory.exact else "(MPS 采样下界)",
)
return meter.avg
def save_backbone(self, epoch: int) -> Path:
"""只存骨架权重,不含投影头——投影头训练成对增强不变,会丢掉下游有用的信息。"""
path = self.work_dir / "simclr_backbone.pth"
torch.save({"model": self.backbone.state_dict(), "epochs": epoch + 1}, path)
return path
def fit(self) -> Path:
"""训练并保存骨架权重(不含投影头),返回权重路径。"""
"""训练并保存骨架权重(不含投影头),返回权重路径。
每 `save_interval` 个 epoch 落一次盘,中途被 Ctrl-C 或 kill 掉也能拿到当前权重继续用。
"""
if len(self.dataloader) < 2:
raise ValueError("SimCLR 需要多个 batch,请调小 batch_size 或增加样本")
for epoch in range(self.epochs):
self.train_one_epoch(epoch)
if self.scheduler is not None:
self.scheduler.step()
path = self.work_dir / "simclr_backbone.pth"
torch.save({"model": self.backbone.state_dict(), "epochs": self.epochs}, path)
try:
for epoch in range(self.epochs):
self.train_one_epoch(epoch)
if self.scheduler is not None:
self.scheduler.step()
if self.save_interval and (epoch + 1) % self.save_interval == 0:
self.save_backbone(epoch)
logger.info("已保存中间权重(epoch %d): %s", epoch, path)
except KeyboardInterrupt:
self.save_backbone(epoch)
logger.warning("训练被中断,已保存 epoch %d 的权重: %s", epoch, path)
return path
self.save_backbone(self.epochs - 1)
logger.info("自监督骨架已保存: %s", path)
return path
+16 -12
View File
@@ -1,7 +1,4 @@
"""训练循环。
注意:训练入口必须先通过 `utils.device.require_cuda()` 的 GPU 检查。
"""
"""训练循环。支持 cuda / mps / cpu 三种后端。"""
from __future__ import annotations
@@ -21,6 +18,7 @@ from engine.losses import DetectionLoss
from engine.metrics import AverageMeter
from engine.optim import build_optimizer, build_scheduler
from engine.visualizer import SwanLabLogger
from utils.device import PeakMemory, resolve_amp
from utils.labels import MULTILABEL
from utils.logging import get_logger
@@ -66,8 +64,10 @@ class Trainer:
self.optimizer = build_optimizer(model, self.cfg.get("optimizer", {"name": "adamw", "lr": 1e-3}))
self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs)
self.amp = bool(self.cfg.get("amp", True)) and device.type == "cuda"
self.scaler = torch.amp.GradScaler("cuda", enabled=self.amp)
# amp: auto 时随后端(CUDA 开、MPS 关——MPS 上 autocast 实测比 fp32 慢)
self.amp = resolve_amp(self.cfg.get("amp", "auto"), device)
self.scaler = torch.amp.GradScaler(device.type, enabled=self.amp)
self.peak_memory = PeakMemory(device)
self.evaluator = (
Evaluator(
@@ -97,7 +97,7 @@ class Trainer:
for step, batch in enumerate(self.train_loader):
batch = _to_device(batch, self.device)
with torch.amp.autocast("cuda", enabled=self.amp):
with torch.amp.autocast(self.device.type, enabled=self.amp):
outputs = self.model(batch["images"])
losses = self.criterion(outputs, batch)
@@ -135,14 +135,18 @@ class Trainer:
return {key: meter.avg for key, meter in meters.items()}
def _gpu_memory(self) -> dict[str, float]:
"""本轮训练的显存峰值(GB);推理端的详细分析见 scripts/profile_inference.py。"""
if self.device.type != "cuda":
"""本轮训练的显存峰值(GB);推理端的详细分析见 scripts/profile_inference.py。
CUDA 读内置峰值计数器;MPS 没有峰值 API,用逐 step 采样的最大值(是下界);
CPU 不统计。
"""
if self.device.type == "cpu":
return {}
stats = {
"gpu_mem_allocated_gb": torch.cuda.max_memory_allocated(self.device) / 1024**3,
"gpu_mem_reserved_gb": torch.cuda.max_memory_reserved(self.device) / 1024**3,
"gpu_mem_allocated_gb": self.peak_memory.peak_mb / 1024,
"gpu_mem_reserved_gb": self.peak_memory.reserved_mb / 1024,
}
torch.cuda.reset_peak_memory_stats(self.device)
self.peak_memory.reset()
return stats
def fit(self) -> float:
+9 -3
View File
@@ -4,7 +4,7 @@
用法:
python eval.py --config configs/eval/baseline.yaml --checkpoint runs/baseline/best.pth
与训练不同,评估允许在 CPU 上运行(`--set device=cpu`),便于部署前回归。
设备同样支持 auto / cuda / mps / cpu;部署前回归可以显式 `--device cpu`。
"""
from __future__ import annotations
@@ -18,7 +18,7 @@ from datasets import build_dataloader, build_dataset, load_classes
from engine import Evaluator
from models import build_model
from utils import check_label_mode, load_config, setup_logging
from utils.device import resolve_device
from utils.device import DeviceNotAvailableError, describe_device, resolve_device
def parse_args() -> argparse.Namespace:
@@ -27,6 +27,7 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--checkpoint", default=None, help="覆盖配置中的 checkpoint")
parser.add_argument("--split", default=None, help="评估哪个 data 划分,默认取配置的 eval.split")
parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE")
parser.add_argument("--device", default=None, help="覆盖配置中的 device:auto / cuda / mps / cpu")
parser.add_argument("--out", default=None, help="把指标写入 json 文件")
return parser.parse_args()
@@ -45,7 +46,12 @@ def main() -> int:
logger.error("必须通过 --checkpoint 或配置中的 checkpoint 指定权重")
return 1
device = resolve_device(str(cfg.get("device", "cuda")))
try:
device = resolve_device(args.device or cfg.get("device", "auto"))
except DeviceNotAvailableError as exc:
logger.error("%s", exc)
return 1
logger.info("评估设备: %s", describe_device(device))
classes = load_classes(cfg["classes"])
# label_mode 由 yaml 决定,模型与数据集必须一致,这里统一下发。
+2 -2
View File
@@ -1,7 +1,7 @@
# 无监督聚类流水线
从 `data/film` 的原始图片到可人工核对的分组结果:六个阶段的数据流、张量形状与每步的数学操作。
虚线是可选的 SimCLR 自监督分支(`method: simclr`,需要 Nvidia GPU)。
虚线是可选的 SimCLR 自监督分支(`method: simclr`,CUDA 与 Apple MPS 均可)。
```mermaid
flowchart TB
@@ -67,7 +67,7 @@ flowchart TB
summ["summary.json<br/>簇大小 / 轮廓 / 稳定性"]:::out
end
simclr["可选:SimCLR 自监督<br/>method = simclr<br/>需 Nvidia GPU"]:::feat
simclr["可选:SimCLR 自监督<br/>method = simclr<br/>cuda / mps 均可"]:::feat
scan --> PRE
PRE -->|"x̃ᵢ"| FEAT
+1 -1
View File
@@ -61,7 +61,7 @@ flowchart TB
summ["summary.json<br/>簇大小 / 轮廓 / 稳定性"]:::out
end
simclr["可选:SimCLR 自监督<br/>method = simclr<br/>需 Nvidia GPU"]:::feat
simclr["可选:SimCLR 自监督<br/>method = simclr<br/>cuda / mps 均可"]:::feat
scan --> PRE
PRE -->|"x̃ᵢ"| FEAT
Binary file not shown.

Before

Width:  |  Height:  |  Size: 193 KiB

After

Width:  |  Height:  |  Size: 193 KiB

+28 -13
View File
@@ -1,7 +1,7 @@
#!/usr/bin/env python
"""环境自检:确认依赖与 Nvidia GPU 是否就绪。
"""环境自检:确认依赖与可用的计算后端。
python scripts/check_env.py
python scripts/check_env.py
"""
from __future__ import annotations
@@ -11,29 +11,44 @@ from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from utils.device import NoNvidiaGpuError, cuda_devices, require_cuda # noqa: E402
from utils.device import ( # noqa: E402
available_backends,
default_amp,
describe_device,
memory_stats,
resolve_device,
)
def main() -> int:
print(f"Python: {sys.version.split()[0]}")
try:
import torch
print(f"PyTorch: {torch.__version__} (CUDA build: {torch.version.cuda})")
except ImportError:
print("PyTorch: 未安装 —— 先执行 `uv sync`")
return 1
devices = cuda_devices()
print(f"Nvidia GPU: {devices if devices else '无'}")
print(f"PyTorch: {torch.__version__} (CUDA 构建: {torch.version.cuda})")
print(f"可用后端(按优先级): {available_backends()}")
try:
require_cuda()
except NoNvidiaGpuError as exc:
print(f"训练不可用: {exc}")
print("(评估/导出仍可在 CPU 上进行)")
print("\n后端明细:")
print(f" cuda: available={torch.cuda.is_available()} device_count={torch.cuda.device_count()}")
print(
f" mps : available={torch.backends.mps.is_available()} "
f"built={torch.backends.mps.is_built()}"
)
device = resolve_device("auto")
print(f"\ndevice=auto 会选中: {describe_device(device)}")
print(f" AMP 默认: {'开' if default_amp(device) else '关'}")
stats = memory_stats(device)
if stats:
print(" " + " ".join(f"{k}={v:.0f}" for k, v in stats.items()))
if device.type == "cpu":
print("\n⚠️ 只有 CPU 可用,训练会非常慢(实测比 MPS 慢约 30 倍),建议换有 GPU 的机器。")
return 1
print("训练环境就绪。")
print("\n训练环境就绪。")
return 0
+13 -13
View File
@@ -4,7 +4,7 @@
# 默认 4 类,ImageNet 预训练特征(不训练,CPU 可跑)
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml
# 改簇数 / 换成在本地数据上自监督训练的特征(需要 Nvidia GPU)
# 改簇数 / 换成在本地数据上自监督训练的特征(Mac 上自动走 MPS)
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
--set num_clusters=6 method=simclr
@@ -40,7 +40,11 @@ from engine.cluster import ( # noqa: E402
)
from models.registry import BACKBONES # noqa: E402
from utils import load_config, seed_everything, setup_logging # noqa: E402
from utils.device import NoNvidiaGpuError, require_cuda, resolve_device # noqa: E402
from utils.device import ( # noqa: E402
DeviceNotAvailableError,
describe_device,
resolve_device,
)
logger = None
@@ -49,7 +53,7 @@ def build_encoder(cfg, device, work_dir):
"""按 `method` 拿一个特征提取器。
pretrained: ImageNet 预训练卷积骨架,纯推理,CPU 也能跑。用顶层的 `backbone`。
simclr: 在本目录数据上自监督训练一个项目自己的卷积骨架——属于训练,先过 GPU 守卫。
simclr: 在本目录数据上自监督训练一个项目自己的卷积骨架。
用 `simclr.backbone`,与 pretrained 的骨架配置各自独立:两者的字段完全不同
(arch/pretrained vs widths/depths),共用一个块会在 `_base_` 合并时串味。
"""
@@ -62,7 +66,6 @@ def build_encoder(cfg, device, work_dir):
if method != "simclr":
raise ValueError(f"未知 method: {method},可选 pretrained / simclr")
require_cuda() # 自监督预训练是训练,禁止在无 Nvidia GPU 的设备上跑
from torch.utils.data import DataLoader
from engine.simclr import SimCLRTrainer
@@ -200,6 +203,7 @@ def main() -> int:
parser.add_argument("--root", default=None, help="覆盖配置里的图片根目录")
parser.add_argument("--num-clusters", type=int, default=None, help="覆盖配置里的簇数")
parser.add_argument("--work-dir", default=None)
parser.add_argument("--device", default=None, help="覆盖配置中的 device:auto / cuda / mps / cpu")
parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE")
parser.add_argument("--save-features", action="store_true", help="额外保存 embedding 供复用")
parser.add_argument(
@@ -220,26 +224,22 @@ def main() -> int:
num_clusters = int(cfg.get("num_clusters", 4))
try:
device = resolve_device(str(cfg.get("device", "cuda")))
except NoNvidiaGpuError as exc:
device = resolve_device(args.device or cfg.get("device", "auto"))
except DeviceNotAvailableError as exc:
logger.error("%s", exc)
return 1
dataset = build_dataset(dict(cfg["dataset"]))
logger.info(
"扫描到 %d 张图,来自 %d 个目录;簇数 %d,method=%s,device=%s",
"扫描到 %d 张图,来自 %d 个目录;簇数 %d,method=%s,设备 %s",
len(dataset),
len({dataset.source_dir(path) for path in dataset.images}),
num_clusters,
cfg.get("method", "pretrained"),
device,
describe_device(device),
)
try:
encoder = build_encoder(cfg, device, work_dir)
except NoNvidiaGpuError as exc:
logger.error("%s", exc)
return 1
encoder = build_encoder(cfg, device, work_dir)
from torch.utils.data import DataLoader
+31 -17
View File
@@ -31,6 +31,12 @@ from torch.profiler import ProfilerActivity, profile, record_function # noqa: E
from datasets import load_classes # noqa: E402
from models import build_model, decode_heatmap # noqa: E402
from utils import load_config, setup_logging # noqa: E402
from utils.device import ( # noqa: E402
DeviceNotAvailableError,
PeakMemory,
resolve_device,
synchronize,
)
logger = setup_logging()
@@ -45,7 +51,7 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--image-size", type=int, nargs=2, default=None, help="默认取配置的 image_size")
parser.add_argument("--iters", type=int, default=50, help="计时轮数")
parser.add_argument("--warmup", type=int, default=10, help="预热轮数(不计入统计)")
parser.add_argument("--device", default=None, help="默认有 GPU 用 cuda,否则 cpu")
parser.add_argument("--device", default=None, help="auto / cuda / mps / cpu,默认 auto")
parser.add_argument("--amp", action="store_true", help="用 float16 autocast 推理")
parser.add_argument("--postprocess", action="store_true", help="把 heatmap 解码也计入")
parser.add_argument("--rows", type=int, default=15, help="算子排行显示条数")
@@ -73,7 +79,7 @@ def build(args: argparse.Namespace, device: torch.device) -> torch.nn.Module:
@torch.no_grad()
def run_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> None:
device_type = images.device.type
with torch.autocast(device_type, dtype=torch.float16, enabled=args.amp and device_type == "cuda"):
with torch.autocast(device_type, dtype=torch.float16, enabled=args.amp and device_type != "cpu"):
with record_function("forward"):
outputs = model(images)
if args.postprocess:
@@ -89,19 +95,17 @@ def run_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namesp
def measure(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> dict:
"""预热后计时,并统计 CUDA 显存峰值。"""
device = images.device
is_cuda = device.type == "cuda"
peak = PeakMemory(device)
for _ in range(args.warmup):
run_once(model, images, args)
if is_cuda:
torch.cuda.synchronize(device)
torch.cuda.reset_peak_memory_stats(device)
synchronize(device)
peak.reset()
start = time.perf_counter()
for _ in range(args.iters):
run_once(model, images, args)
if is_cuda:
torch.cuda.synchronize(device)
synchronize(device)
elapsed = time.perf_counter() - start
batch = images.shape[0]
@@ -110,11 +114,12 @@ def measure(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespa
"latency_ms": elapsed / args.iters * 1000,
"throughput_fps": args.iters * batch / elapsed,
}
if is_cuda:
if device.type != "cpu":
result.update(
{
"peak_allocated_mb": torch.cuda.max_memory_allocated(device) / 1024**2,
"peak_reserved_mb": torch.cuda.max_memory_reserved(device) / 1024**2,
"peak_allocated_mb": peak.peak_mb,
"peak_reserved_mb": peak.reserved_mb,
"peak_is_exact": peak.exact,
"weights_mb": sum(p.numel() * p.element_size() for p in model.parameters()) / 1024**2,
}
)
@@ -124,13 +129,13 @@ def measure(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespa
def profile_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Namespace) -> None:
"""算子级 profiling:按显存占用与耗时各排一次序。"""
is_cuda = images.device.type == "cuda"
# torch.profiler 没有 MPS 的 activity,Metal 侧的 kernel 抓不到,只能看 CPU 侧下发耗时
activities = [ProfilerActivity.CPU] + ([ProfilerActivity.CUDA] if is_cuda else [])
with profile(activities=activities, profile_memory=True, record_shapes=True) as prof:
for _ in range(max(args.iters // 5, 1)):
run_once(model, images, args)
if is_cuda:
torch.cuda.synchronize(images.device)
synchronize(images.device)
memory_key = "self_cuda_memory_usage" if is_cuda else "self_cpu_memory_usage"
time_key = "self_cuda_time_total" if is_cuda else "self_cpu_time_total"
@@ -148,9 +153,18 @@ def profile_once(model: torch.nn.Module, images: torch.Tensor, args: argparse.Na
def main() -> int:
args = parse_args()
device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu"))
if device.type != "cuda":
logger.warning("当前在 %s 上运行:没有显存数据,只能看耗时与 CPU 内存。", device)
try:
device = resolve_device(args.device or "auto")
except DeviceNotAvailableError as exc:
logger.error("%s", exc)
return 1
if device.type == "cpu":
logger.warning("当前在 CPU 上运行:没有显存数据,只能看耗时与 CPU 内存。")
elif device.type == "mps":
logger.warning(
"MPS 上 torch.profiler 抓不到 Metal kernel(没有对应的 ProfilerActivity),"
"算子明细只反映 CPU 侧下发耗时;显存用 torch.mps 采样,峰值是下界。"
)
model = build(args, device)
cfg = model.cfg
@@ -158,7 +172,7 @@ def main() -> int:
channels = 1 if cfg.get("gray") else 3
logger.info(
"模型 %s | 参数量 %.2fM | 输出 stride %d | 输入 %dx%dx%d | device %s | amp=%s",
"模型 %s | 参数量 %.2fM | 输出 stride %d | 输入 %dx%dx%d | 设备 %s | amp=%s",
cfg["model"]["name"],
sum(p.numel() for p in model.parameters()) / 1e6,
model.out_stride,
+112
View File
@@ -0,0 +1,112 @@
"""设备解析:cuda / mps / cpu 三后端的选择、AMP 默认值与显存统计。"""
import pytest
torch = pytest.importorskip("torch")
from utils.device import ( # noqa: E402
BACKEND_PRIORITY,
DeviceNotAvailableError,
PeakMemory,
amp_supported,
available_backends,
default_amp,
describe_device,
memory_stats,
pin_memory_supported,
resolve_amp,
resolve_device,
synchronize,
)
def test_available_backends_is_ordered_and_always_has_cpu():
backends = available_backends()
assert backends[-1] == "cpu"
assert backends == [b for b in BACKEND_PRIORITY if b in backends]
assert len(set(backends)) == len(backends)
def test_auto_picks_the_highest_priority_available_backend():
device = resolve_device("auto")
assert device.type == available_backends()[0]
assert resolve_device("") == device # 空串等同 auto
def test_cpu_is_always_resolvable():
assert resolve_device("cpu").type == "cpu"
def test_explicit_backend_matches_hardware_or_raises():
"""显式请求不可用的后端必须报错,不能悄悄退回 CPU——那会让训练慢几十倍还不吭声。"""
for backend in ("cuda", "mps"):
if backend in available_backends():
assert resolve_device(backend).type == backend
else:
with pytest.raises(DeviceNotAvailableError, match=backend):
resolve_device(backend)
def test_unknown_device_raises():
with pytest.raises(DeviceNotAvailableError, match="未知设备"):
resolve_device("tpu")
def test_amp_defaults_follow_the_backend():
# CUDA 有 tensor core,默认开;MPS 上 autocast 实测更慢,默认关;CPU 不支持
assert default_amp("cuda") is True
assert default_amp("mps") is False
assert default_amp("cpu") is False
assert amp_supported("cuda") and amp_supported("mps")
assert not amp_supported("cpu")
def test_resolve_amp_handles_auto_and_explicit():
assert resolve_amp("auto", "cuda") is True
assert resolve_amp("auto", "mps") is False
assert resolve_amp(None, "mps") is False
assert resolve_amp(True, "mps") is True # 显式打开,MPS 支持 autocast
assert resolve_amp(True, "cpu") is False # CPU 不支持,显式打开也无效
assert resolve_amp(False, "cuda") is False
def test_pin_memory_only_for_cuda():
assert pin_memory_supported("cuda")
assert not pin_memory_supported("mps")
assert not pin_memory_supported("cpu")
def test_describe_device_is_readable():
assert describe_device("cpu") == "CPU"
assert "MPS" in describe_device("mps")
def test_memory_stats_shape_per_backend():
assert memory_stats("cpu") == {}
stats = memory_stats(resolve_device("auto"))
device_type = resolve_device("auto").type
if device_type == "cuda":
assert {"allocated_mb", "reserved_mb", "peak_allocated_mb"} <= set(stats)
elif device_type == "mps":
assert {"allocated_mb", "driver_allocated_mb"} <= set(stats)
def test_peak_memory_tracks_growth_on_accelerators():
device = resolve_device("auto")
peak = PeakMemory(device)
if device.type == "cpu":
assert peak.peak_mb == 0.0 and peak.reserved_mb == 0.0
return
# MPS 没有峰值 API,靠采样,所以峰值是下界(exact=False)
assert peak.exact == (device.type != "mps")
buffer = torch.empty(8 * 1024 * 1024 // 4, device=device) # 约 8 MB
peak.update()
synchronize(device)
assert peak.peak_mb >= 4.0
del buffer
def test_synchronize_is_a_noop_on_cpu():
synchronize("cpu") # 不应抛异常
-26
View File
@@ -1,26 +0,0 @@
"""GPU 守卫:项目禁止在没有 Nvidia GPU 的设备上训练。"""
import pytest
from utils.device import NoNvidiaGpuError, cuda_devices, require_cuda
def test_require_cuda_matches_hardware():
devices = cuda_devices()
if devices:
assert require_cuda() == devices
else:
with pytest.raises(NoNvidiaGpuError):
require_cuda()
def test_resolve_device_rejects_cuda_without_gpu():
torch = pytest.importorskip("torch")
from utils.device import resolve_device
if torch.cuda.is_available():
assert resolve_device("cuda").type == "cuda"
else:
with pytest.raises(NoNvidiaGpuError):
resolve_device("cuda")
assert resolve_device("cpu").type == "cpu"
+49 -10
View File
@@ -13,7 +13,6 @@ from PIL import Image # noqa: E402
from engine.simclr import SimCLRTrainer, nt_xent_loss # noqa: E402
from models.heads.projection import ProjectionHead # noqa: E402
from models.registry import BACKBONES # noqa: E402
from utils.device import NoNvidiaGpuError # noqa: E402
def test_nt_xent_is_minimal_when_views_agree():
@@ -105,12 +104,46 @@ def test_simclr_trainer_runs_and_saves_backbone_only(tiny_image_dir, tmp_path):
backbone.load_state_dict(state["model"]) # 能被干净地加载回去
def test_simclr_saves_intermediate_checkpoints(tiny_image_dir, tmp_path):
"""中途落盘:跑到一半被打断也要能拿到权重(此前只在全部 epoch 跑完后才存)。"""
from torch.utils.data import DataLoader
from datasets.unlabeled import UnlabeledImageDataset, unlabeled_collate
dataset = UnlabeledImageDataset(
tiny_image_dir,
transforms=[{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}],
views=2,
)
loader = DataLoader(dataset, batch_size=4, shuffle=True, drop_last=True, collate_fn=unlabeled_collate)
backbone = BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]})
work_dir = tmp_path / "run"
trainer = SimCLRTrainer(
backbone,
loader,
{"epochs": 4, "amp": False, "log_interval": 100, "save_interval": 2,
"projection": {"hidden_channels": 16, "out_channels": 8}},
torch.device("cpu"),
work_dir,
)
# 手动走前两个 epoch,模拟"训练中途"
trainer.train_one_epoch(0)
trainer.train_one_epoch(1)
path = trainer.save_backbone(1)
assert path.exists()
assert torch.load(path, map_location="cpu", weights_only=False)["epochs"] == 2
trainer.fit() # 完整跑完后 epochs 应更新
assert torch.load(path, map_location="cpu", weights_only=False)["epochs"] == 4
def test_cluster_script_simclr_branch_end_to_end(tiny_image_dir, tmp_path, monkeypatch):
"""method=simclr 全流程:训练 -> 存权重 -> 抽特征 -> 聚类 -> 产出。"""
sys.path.insert(0, "scripts")
import cluster_unlabeled as cu
monkeypatch.setattr(cu, "require_cuda", lambda: ["FAKE-GPU"])
work_dir = tmp_path / "run"
monkeypatch.setattr(
sys, "argv",
@@ -130,19 +163,25 @@ def test_cluster_script_simclr_branch_end_to_end(tiny_image_dir, tmp_path, monke
assert (work_dir / "clusters" / "cluster_00").is_dir()
def test_simclr_branch_is_blocked_without_gpu(tiny_image_dir, tmp_path, monkeypatch):
"""无 Nvidia GPU 时 simclr 分支必须被守卫拦下(pretrained 分支不受影响)。"""
def test_simclr_runs_on_the_resolved_device(tiny_image_dir, tmp_path, monkeypatch):
"""method=simclr 在本机自动选中的后端上(MPS / CUDA / CPU)都应能跑通。"""
sys.path.insert(0, "scripts")
import cluster_unlabeled as cu
def refuse():
raise NoNvidiaGpuError("没有 Nvidia GPU")
from utils.device import describe_device, resolve_device
monkeypatch.setattr(cu, "require_cuda", refuse)
device = resolve_device("auto")
work_dir = tmp_path / "auto_run"
monkeypatch.setattr(
sys, "argv",
["cluster_unlabeled.py", "--config", "configs/cluster/film.yaml",
"--work-dir", str(tmp_path / "run"), "--root", str(tiny_image_dir),
"--set", "device=cpu", "method=simclr", "loader.num_workers=0"],
"--work-dir", str(work_dir), "--root", str(tiny_image_dir), "--set",
"method=simclr", "loader.num_workers=0", "loader.batch_size=4",
"simclr.epochs=1", "simclr.batch_size=4",
"simclr.backbone.widths=[8,16]", "simclr.backbone.depths=[1,1]",
"simclr.projection.hidden_channels=16", "simclr.projection.out_channels=8",
"num_clusters=2", "n_init=2"],
)
assert cu.main() == 1
assert cu.main() == 0, f"在 {describe_device(device)} 上跑失败"
assert (work_dir / "simclr_backbone.pth").exists()
+26 -11
View File
@@ -1,7 +1,7 @@
"""train.py / eval.py 的端到端冒烟:YOLO 数据目录 -> 训练 -> 评估。
覆盖入口脚本的接线(classes 解析、label_mode 下发、GPU 守卫),
所以在 CPU 上跑,`require_cuda` 用 monkeypatch 放行。
覆盖入口脚本的接线:classes 解析、label_mode 下发、设备解析。固定用 CPU 跑,
免得测试受机器上有没有 GPU 影响。
"""
import sys
@@ -19,7 +19,6 @@ from PIL import Image # noqa: E402
import eval as eval_mod # noqa: E402
import train as train_mod # noqa: E402
from datasets import write_data_yaml # noqa: E402
from utils.device import NoNvidiaGpuError # noqa: E402
CLASSES = ["scratch", "bubble"]
TRANSFORMS = [
@@ -85,7 +84,6 @@ def run(module, argv, monkeypatch):
def test_train_then_eval(project, monkeypatch, tmp_path):
config_path, root = project
work_dir = tmp_path / "runs"
monkeypatch.setattr(train_mod, "require_cuda", lambda: ["FAKE-GPU"])
assert run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)],
monkeypatch) == 0
@@ -115,7 +113,6 @@ def test_classes_from_data_yaml_sizes_the_heads(project, monkeypatch, tmp_path):
return real_build(cfg)
monkeypatch.setattr(train_mod, "build_model", spy)
monkeypatch.setattr(train_mod, "require_cuda", lambda: ["FAKE-GPU"])
run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)], monkeypatch)
@@ -123,15 +120,33 @@ def test_classes_from_data_yaml_sizes_the_heads(project, monkeypatch, tmp_path):
assert captured["cfg"]["label_mode"] == "multilabel"
def test_training_is_blocked_without_nvidia_gpu(project, monkeypatch, tmp_path):
config_path, _ = project
def test_unavailable_device_fails_loudly(project, monkeypatch, tmp_path):
"""显式请求机器上没有的后端时必须退出,而不是悄悄退回 CPU。"""
from utils.device import available_backends
def refuse():
raise NoNvidiaGpuError("没有 Nvidia GPU")
missing = next((b for b in ("cuda", "mps") if b not in available_backends()), None)
if missing is None:
pytest.skip("这台机器 cuda 与 mps 都可用")
monkeypatch.setattr(train_mod, "require_cuda", refuse)
assert run(
train_mod,
["train.py", "--config", str(config_path), "--work-dir", str(tmp_path / "runs3")],
["train.py", "--config", str(config_path_of(project)), "--work-dir", str(tmp_path / "runs3"),
"--device", missing],
monkeypatch,
) == 1
def test_device_flag_overrides_config(project, monkeypatch, tmp_path):
"""--device 优先于配置里的 device。"""
config_path, _ = project
work_dir = tmp_path / "runs4"
assert run(
train_mod,
["train.py", "--config", str(config_path), "--work-dir", str(work_dir), "--device", "cpu"],
monkeypatch,
) == 0
assert (work_dir / "best.pth").exists()
def config_path_of(project):
return project[0]
+8 -8
View File
@@ -4,8 +4,10 @@
用法:
python train.py --config configs/train/baseline.yaml
python train.py --config configs/train/baseline.yaml --set train.epochs=10 train.optimizer.lr=5e-4
python train.py --config configs/train/baseline.yaml --device mps
硬性约束:没有 Nvidia GPU 时直接退出(见 CLAUDE.MD / utils/device.py)。
设备由配置的 `device` 决定,默认 `auto`:按 cuda > mps > cpu 挑最快的可用后端。
显式指定但不可用时直接报错,不会悄悄退回 CPU。
"""
from __future__ import annotations
@@ -14,13 +16,11 @@ import argparse
import sys
from pathlib import Path
import torch
from datasets import build_dataloader, build_dataset, load_classes
from engine import Trainer
from models import build_model
from utils import check_label_mode, load_config, seed_everything, setup_logging
from utils.device import NoNvidiaGpuError, require_cuda
from utils.device import DeviceNotAvailableError, describe_device, resolve_device
def parse_args() -> argparse.Namespace:
@@ -28,6 +28,7 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--config", required=True, help="yaml 配置路径")
parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE", help="覆盖配置项")
parser.add_argument("--work-dir", default=None, help="覆盖配置中的 work_dir")
parser.add_argument("--device", default=None, help="覆盖配置中的 device:auto / cuda / mps / cpu")
parser.add_argument("--resume", default=None, help="从检查点恢复训练")
return parser.parse_args()
@@ -41,14 +42,13 @@ def main() -> int:
logger = setup_logging(work_dir / "train.log")
try:
gpus = require_cuda()
except NoNvidiaGpuError as exc:
device = resolve_device(args.device or cfg.get("device", "auto"))
except DeviceNotAvailableError as exc:
logger.error("%s", exc)
return 1
logger.info("检测到 GPU: %s", ", ".join(gpus))
logger.info("训练设备: %s", describe_device(device))
seed_everything(int(cfg.get("seed", 42)), bool(cfg.get("deterministic", False)))
device = torch.device(cfg.get("device", "cuda"))
classes = load_classes(cfg["classes"])
# label_mode 由 yaml 决定,模型与数据集必须一致,这里统一下发。
+15 -2
View File
@@ -1,7 +1,15 @@
"""通用工具:配置、注册表、日志、随机种子、设备检查。"""
from utils.config import Config, load_config
from utils.device import require_cuda, resolve_device
from utils.device import (
DeviceNotAvailableError,
PeakMemory,
available_backends,
describe_device,
memory_stats,
resolve_amp,
resolve_device,
)
from utils.labels import LABEL_MODES, MULTICLASS, MULTILABEL, check_label_mode
from utils.logging import get_logger, setup_logging
from utils.registry import Registry
@@ -10,8 +18,13 @@ from utils.seed import seed_everything
__all__ = [
"Config",
"load_config",
"require_cuda",
"resolve_device",
"resolve_amp",
"describe_device",
"available_backends",
"memory_stats",
"PeakMemory",
"DeviceNotAvailableError",
"LABEL_MODES",
"MULTILABEL",
"MULTICLASS",
+174 -39
View File
@@ -1,54 +1,189 @@
"""设备检查。
"""设备解析与能力查询:统一处理 cuda / mps / cpu 三种后端。
项目硬性约定(见 CLAUDE.MD):禁止在没有 Nvidia GPU 的设备上训练。
所有训练入口必须先调用 `require_cuda()`,检查不通过直接抛错退出。
设计要点:
- `resolve_device("auto")` 按 cuda > mps > cpu 的优先级自动挑最快的后端;
- 显式指定但不可用时直接报错,而不是悄悄退回 CPU——训练慢 30 倍却不吭声是最坏的情况;
- AMP 的默认值随后端而变:CUDA 上开(有 tensor core),MPS 上关
(实测 tiny_conv @128² 每步 fp32 76 ms、fp16 90 ms,autocast 只剩转换开销反而更慢);
- 显存统计跨后端统一:CUDA 有内置峰值计数器,MPS 没有,靠采样取最大值。
"""
from __future__ import annotations
import torch
class NoNvidiaGpuError(RuntimeError):
"""当前设备不具备可用的 Nvidia GPU。"""
#: 自动选择时的优先级
BACKEND_PRIORITY = ("cuda", "mps", "cpu")
def cuda_devices() -> list[str]:
"""返回可用的 Nvidia GPU 名称列表;没有 torch 或没有 GPU 时返回空列表。"""
try:
import torch
except ImportError:
return []
if not torch.cuda.is_available() or torch.version.cuda is None:
return []
return [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]
class DeviceNotAvailableError(RuntimeError):
"""请求的设备在当前机器上不可用。"""
def require_cuda() -> list[str]:
"""训练前的守卫;无 Nvidia GPU 时抛出 `NoNvidiaGpuError`。"""
try:
import torch
except ImportError as exc:
raise NoNvidiaGpuError("未安装 PyTorch,无法确认 Nvidia GPU,训练已中止。") from exc
if torch.version.cuda is None:
raise NoNvidiaGpuError(
f"当前 PyTorch ({torch.__version__}) 不是 CUDA 构建(可能是 CPU/ROCm/MPS 版本),"
"本项目禁止在非 Nvidia GPU 设备上训练。"
)
if not torch.cuda.is_available() or torch.cuda.device_count() == 0:
raise NoNvidiaGpuError(
"未检测到可用的 Nvidia GPU,本项目禁止在无 GPU 的设备上训练。"
"评估/推理请使用 eval.py 并显式指定 device=cpu。"
)
return cuda_devices()
def _cuda_ready() -> bool:
return torch.cuda.is_available() and torch.cuda.device_count() > 0
def resolve_device(device: str = "cuda"):
"""把配置中的 device 字符串解析为 `torch.device`。
def _mps_ready() -> bool:
return torch.backends.mps.is_available() and torch.backends.mps.is_built()
仅用于评估/推理等允许 CPU 的场景;训练路径请先走 `require_cuda()`。
def available_backends() -> list[str]:
"""按优先级返回当前可用的后端,末尾必定有 cpu。"""
backends = []
if _cuda_ready():
backends.append("cuda")
if _mps_ready():
backends.append("mps")
backends.append("cpu")
return backends
def resolve_device(spec: str = "auto") -> torch.device:
"""把配置里的 device 字符串解析成 `torch.device`。
Args:
spec: "auto" 自动挑最优后端;也可写 "cuda" / "cuda:1" / "mps" / "cpu"。
Raises:
DeviceNotAvailableError: 显式请求的后端不可用。
"""
import torch
spec = str(spec).strip().lower()
if spec in ("auto", ""):
return torch.device(available_backends()[0])
if device.startswith("cuda") and not torch.cuda.is_available():
raise NoNvidiaGpuError(f"请求了 {device},但当前设备没有可用的 Nvidia GPU。")
return torch.device(device)
backend = spec.split(":", 1)[0]
if backend == "cuda" and not _cuda_ready():
raise DeviceNotAvailableError(
f"请求了 {spec},但当前没有可用的 Nvidia GPU"
f"(PyTorch {torch.__version__},CUDA 构建:{torch.version.cuda})。"
f"可用后端:{available_backends()},或用 device=auto 自动选择。"
)
if backend == "mps" and not _mps_ready():
raise DeviceNotAvailableError(
f"请求了 {spec},但当前 PyTorch 不支持 MPS"
f"(available={torch.backends.mps.is_available()}, built={torch.backends.mps.is_built()})。"
f"可用后端:{available_backends()}。"
)
if backend not in ("cuda", "mps", "cpu"):
raise DeviceNotAvailableError(f"未知设备 {spec},支持 cuda / mps / cpu / auto。")
return torch.device(spec)
def describe_device(device: torch.device | str) -> str:
"""给日志用的人类可读描述。"""
device = torch.device(device)
if device.type == "cuda":
index = device.index or 0
return f"{torch.cuda.get_device_name(index)} (cuda:{index})"
if device.type == "mps":
return "Apple Silicon GPU (Metal / MPS)"
return "CPU"
def amp_supported(device: torch.device | str) -> bool:
"""该后端是否支持 autocast 混合精度。"""
return torch.device(device).type in ("cuda", "mps")
def default_amp(device: torch.device | str) -> bool:
"""AMP 的默认开关:只有 CUDA 默认开。
MPS 支持 autocast 且数值正常,但实测比 fp32 更慢,所以默认关闭;
需要省显存时可以在配置里显式 `amp: true` 打开。
"""
return torch.device(device).type == "cuda"
def resolve_amp(setting: bool | str | None, device: torch.device | str) -> bool:
"""解析配置里的 `amp` 字段:`auto`/`None` 随后端,`true`/`false` 显式指定。"""
device = torch.device(device)
if setting is None or (isinstance(setting, str) and setting.lower() == "auto"):
return default_amp(device)
if not bool(setting):
return False
return amp_supported(device)
def pin_memory_supported(device: torch.device | str) -> bool:
"""只有 CUDA 的锁页内存有意义;MPS 上开启会告警且无收益。"""
return torch.device(device).type == "cuda"
def synchronize(device: torch.device | str) -> None:
"""等待该设备上的异步算子完成——计时前后必须调用,否则测到的是下发时间。"""
device = torch.device(device)
if device.type == "cuda":
torch.cuda.synchronize(device)
elif device.type == "mps":
torch.mps.synchronize()
def memory_stats(device: torch.device | str) -> dict[str, float]:
"""统一的显存查询,单位 MB。
CUDA 返回 allocated / reserved / peak;MPS 只有当前分配量与驱动占用
(Metal 没有暴露峰值计数器);CPU 返回空字典。
"""
device = torch.device(device)
mb = 1024.0 * 1024.0
if device.type == "cuda":
return {
"allocated_mb": torch.cuda.memory_allocated(device) / mb,
"reserved_mb": torch.cuda.memory_reserved(device) / mb,
"peak_allocated_mb": torch.cuda.max_memory_allocated(device) / mb,
"peak_reserved_mb": torch.cuda.max_memory_reserved(device) / mb,
}
if device.type == "mps":
return {
"allocated_mb": torch.mps.current_allocated_memory() / mb,
"driver_allocated_mb": torch.mps.driver_allocated_memory() / mb,
}
return {}
class PeakMemory:
"""跨后端的显存峰值统计。
CUDA 有内置峰值计数器,直接读;MPS 没有,只能在 `update()` 时采样取最大值——
因此 MPS 上的峰值是**下界**(采样点之间的瞬时尖峰会被漏掉)。
"""
def __init__(self, device: torch.device | str) -> None:
self.device = torch.device(device)
self._peak = 0.0
self.reset()
def reset(self) -> None:
self._peak = 0.0
if self.device.type == "cuda":
torch.cuda.reset_peak_memory_stats(self.device)
self.update()
def update(self) -> None:
if self.device.type != "mps":
return
self._peak = max(self._peak, torch.mps.current_allocated_memory() / (1024.0 * 1024.0))
@property
def peak_mb(self) -> float:
if self.device.type == "cuda":
return torch.cuda.max_memory_allocated(self.device) / (1024.0 * 1024.0)
if self.device.type == "mps":
self.update()
return self._peak
return 0.0
@property
def reserved_mb(self) -> float:
"""CUDA 是缓存分配器保留量;MPS 是 Metal 驱动占用;CPU 为 0。"""
if self.device.type == "cuda":
return torch.cuda.max_memory_reserved(self.device) / (1024.0 * 1024.0)
if self.device.type == "mps":
return torch.mps.driver_allocated_memory() / (1024.0 * 1024.0)
return 0.0
@property
def exact(self) -> bool:
"""峰值是否精确(MPS 靠采样,只能给下界)。"""
return self.device.type != "mps"