PyTorch工程基座：5分钟启动可复现、可调试、可部署的训练流程

1. 项目概述：这不是“入门教程”，而是一套可即插即用的PyTorch工程基座

“PyTorch Starter Pack”——光看名字，很多人会下意识划走：又一个讲torch.nn.Module继承、DataLoader初始化、optimizer.step()三板斧的速成课？但我在带团队从零搭建CV/NLP小模型项目时发现，真正卡住新手的从来不是“怎么写网络”，而是“写完第一行代码后，接下来该做什么”。比如：训练日志打到哪？验证集指标怎么算才不泄露信息？模型保存要不要带torch.compile优化后的图？torch.backends.cudnn.benchmark = True到底该在哪儿设、什么时候不该设？这些细节没有标准答案，但每错一步，轻则训练结果不可复现，重则GPU显存悄无声息爆掉、凌晨三点还在debugRuntimeError: expected scalar type Float but found Half。

这个Starter Pack，就是我过去三年在工业界落地17个中小型PyTorch项目（涵盖图像分类、时序预测、多模态检索、轻量化部署）后，把反复验证过的“最小可行工程骨架”抽离出来形成的。它不教张量运算原理，不讲反向传播推导，只解决一件事：让你在5分钟内启动一个结构清晰、日志完整、可复现、可调试、可扩展的训练流程。核心关键词是：可复现性、模块化组织、生产级日志、设备无关设计、轻量部署友好。适合两类人：一是刚学完《PyTorch官方教程》第3章、面对真实数据集手足无措的在校生；二是需要快速验证算法想法、又不想被工程细节拖垮进度的算法工程师。它不是替代你思考的黑盒，而是帮你把重复劳动压缩到一行命令就能拉起的基座——就像你不会每次造车都重炼钢铁，但得清楚轮子该装在哪、油箱该加什么标号的油。

我试过把这套结构直接交给实习生，他第一天就跑通了自己收集的花卉图像分类任务，第二天开始调参，第三天把模型转ONNX部署到树莓派上。关键不在代码多炫酷，而在所有“隐性成本”都被提前封进了配置和约定里：随机种子怎么设才真正覆盖所有随机源？DistributedDataParallel的find_unused_parameters为什么默认关？torch.compile在训练/推理阶段的推荐模式差异是什么？这些答案，全藏在Pack的每一处默认值和注释里。它不承诺“零bug”，但承诺“每个bug都有明确归因路径”。

2. 整体架构设计：为什么放弃“脚本式”而选择“模块化+配置驱动”

2.1 传统单文件训练脚本的三大死穴

很多入门教程推崇“一个py文件搞定所有”，比如train.py里堆满if __name__ == '__main__':下的逻辑。实测下来，这种结构在项目超过3个实验、2种数据源、1种模型变体后，就会迅速崩坏。我整理过团队早期的train_v1.py到train_v7.py迭代记录，发现三个高频痛点：

• 复现性灾难：随机种子只设了torch.manual_seed(42)，却忘了numpy.random.seed(42)、random.seed(42)，更没碰torch.cuda.manual_seed_all(42)。结果同一份代码，在A卡上ACC 89.2%，B卡上变成88.7%，排查三天才发现是CUDA版本差异导致cudnn.deterministic行为不一致。

• 参数耦合地狱：学习率、batch_size、warmup步数全写死在代码里。想对比不同lr，就得改代码、git commit、再跑——而不是python train.py --lr 1e-3一键切换。更糟的是，当你要加一个新功能（比如梯度裁剪），得在optimizer.step()前后各插一段逻辑，极易漏掉某处。

• 日志与监控失联：print语句满天飞，但loss曲线画不出来；tensorboard日志路径硬编码，换机器就得改；验证指标只打印不保存，想回溯上周的mAP？抱歉，终端早已滚动消失。

提示：Starter Pack的第一条铁律是——任何可能变化的参数，必须从代码中剥离，进入配置层。这不是为了炫技，而是让“实验管理”从玄学变成可操作动作。

2.2 我们采用的三层架构：config → core → cli

整个Pack按职责严格分层，目录结构如下（精简后）：

stater_pack/ ├── configs/ # 所有可变参数的唯一源头 │ ├── base.yaml # 全局默认（设备、种子、日志路径） │ ├── model/ # 模型相关（arch, hidden_dim, dropout） │ │ └── resnet18.yaml │ ├── data/ # 数据相关（root, img_size, augment） │ │ └── flowers.yaml │ └── train/ # 训练策略（lr, epochs, scheduler） │ └── default.yaml ├── core/ # 纯逻辑，零配置硬编码 │ ├── model/ # 模型定义（nn.Module子类） │ │ └── resnet.py │ ├── data/ # 数据加载（Dataset, DataLoader构建） │ │ └── flowers.py │ ├── trainer.py # 核心训练循环（含DDP支持、compile集成） │ └── utils.py # 工具函数（seed_everything, save_checkpoint） ├── cli.py # 命令行入口（argparse + hydra集成） └── train.py # 用户唯一需执行的脚本（仅10行）

这个设计背后有明确取舍：

• 拒绝Hydra的全部特性：只用其配置合并能力（@hydra.main(config_path="configs", config_name="train")），不用其@hydra.main装饰器的复杂注入机制。原因？Hydra的OmegaConf对象在调试时类型提示混乱，trainer.train()里打个断点，cfg.model的类型是DictConfig而非dict，IDE无法跳转，新人极易懵圈。我们用omegaconf.OmegaConf.to_container(cfg, resolve=True)在入口处转成原生dict，后续全是Python原生类型。

• 配置优先级明确：base.yaml<model/*.yaml<data/*.yaml<train/*.yaml< 命令行参数。例如python train.py model=resnet50 data=cifar10 train=lr_5e-4，会自动合并四层配置，命令行参数最高优先级。这样，你无需复制粘贴yaml文件，一个命令就能组合出新实验。

• core层绝对纯净：core/trainer.py里看不到任何if cfg.model.name == "resnet"的分支判断。模型实例化由core/model/__init__.py的工厂函数完成：

  def build_model(cfg: dict) -> nn.Module: arch = cfg["arch"] if arch == "resnet18": return ResNet18(num_classes=cfg["num_classes"]) elif arch == "vit_tiny": return ViTTiny(num_classes=cfg["num_classes"]) # ... 其他模型

  这样，加新模型只需在core/model/下新增文件+注册工厂函数，不污染训练主逻辑。

2.3 关键设计决策背后的“为什么”

这个架构不是为“看起来高级”而设计，而是为“少踩坑”而存在。当你在深夜调试一个OOM错误时，你会感谢core/trainer.py里那行torch.cuda.empty_cache()被精准放在validate_epoch之后——而不是像某些教程那样，把它丢在train_epoch末尾，导致验证阶段显存反而更高。

3. 核心模块详解：从种子到部署的每一个关键环节

3.1 种子固化：为什么seed_everything(42)还不够

新手常以为设一个torch.manual_seed(42)就万事大吉。但PyTorch生态里，随机性来自至少5个独立源头：

1. CPU随机数生成器：torch.manual_seed()、numpy.random.seed()、random.seed()
2. CUDA随机数生成器：torch.cuda.manual_seed()（单卡）、torch.cuda.manual_seed_all()（多卡）
3. cuDNN卷积算法选择器：torch.backends.cudnn.deterministic = True强制确定性算法，但会牺牲10-15%速度
4. 数据增强随机性：torchvision.transforms.Random*类内部使用torch.Generator，需单独设种子
5. Dataloader worker随机性：DataLoader(num_workers>0)的每个worker有自己的numpy和random状态

Starter Pack的core/utils.py中，seed_everything()函数完整覆盖这五点：

def seed_everything(seed: int): """Set seeds for reproducibility across all random sources.""" import random import numpy as np import torch # 1. CPU sources random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) # 2. CUDA sources (only if available) if torch.cuda.is_available(): torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) # for multi-GPU # 3. cuDNN deterministic torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False # benchmark=True会破坏确定性！ # 4. For torchvision transforms that use Generator # We'll pass generator to transforms in data module # 5. For DataLoader workers, we set worker_init_fn # This is handled in core/data/base.py's get_dataloader()

注意：torch.backends.cudnn.benchmark = False是关键！很多教程说“设True加速”，但它会先尝试多种卷积算法并缓存最优者，这个过程本身是非确定性的。在需要复现的场景，必须关掉。

实操中，我们还做了两件事：

• 在core/data/flowers.py中，FlowersDataset的__getitem__方法接收一个generator参数，并传给torchvision.transforms.RandomHorizontalFlip(p=0.5, generator=generator)
• 在core/data/base.py的get_dataloader()中，worker_init_fn被定义为：

  def worker_init_fn(worker_id): worker_seed = torch.initial_seed() % 2**32 np.random.seed(worker_seed) random.seed(worker_seed)

  这样，每个dataloader worker的随机状态都源于主进程种子，彻底杜绝数据加载阶段的随机漂移。

3.2 数据加载：如何让DataLoader不成为性能瓶颈

数据加载往往是训练Pipeline中最隐蔽的瓶颈。我见过太多案例：GPU利用率长期低于30%，nvidia-smi显示显存占满但计算单元空转——问题就出在DataLoader。

Starter Pack的数据模块设计遵循三个原则：预处理下沉、内存映射优化、异步解耦。

预处理下沉到Dataset

避免在DataLoader的collate_fn里做耗时操作（如torch.stack、torch.cat）。我们在core/data/flowers.py中，FlowersDataset.__getitem__直接返回Tensor：

class FlowersDataset(Dataset): def __init__(self, root: str, split: str, transform: Optional[Callable] = None): self.root = Path(root) self.split = split self.transform = transform # 预加载所有图片路径和标签，避免__getitem__中IO self.samples = self._load_samples() # list of (path, label) def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]: img_path, label = self.samples[idx] # PIL读取 + 转Tensor，全程在CPU内存 img = Image.open(img_path).convert("RGB") if self.transform: img = self.transform(img) # transform已定义为torchvision.transforms.Compose return img, label # img is already torch.Tensor

transform在configs/data/flowers.yaml中定义：

transform: train: - name: "Resize" size: [256, 256] - name: "RandomHorizontalFlip" p: 0.5 - name: "ToTensor" - name: "Normalize" mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225]

core/data/base.py中的build_transform()函数会动态构建Compose，确保所有Random*变换都接收generator参数。

内存映射优化：ImageFoldervs 自定义Dataset

对于标准ImageFolder结构（root/class1/xxx.jpg），我们仍推荐自定义Dataset。因为ImageFolder在__init__中会遍历所有子目录，当数据集达百万级时，os.listdir耗时惊人。我们的_load_samples()方法使用pathlib.Path.rglob("*.jpg")并缓存结果到samples.pkl，首次加载慢，后续秒开。

异步解耦：pin_memory与prefetch_factor

DataLoader的关键参数在configs/train/default.yaml中配置：

dataloader: batch_size: 64 num_workers: 8 pin_memory: true prefetch_factor: 2 # 每个worker预取2个batch persistent_workers: true # worker进程复用，避免反复fork开销

pin_memory=True将CPU Tensor锁页，使GPU能通过DMA直接访问，速度提升约20%。persistent_workers=True在PyTorch 1.7+引入，避免每个epoch重建worker进程的开销（实测epoch间gap从1.2s→0.05s）。

实操心得：num_workers不是越多越好。我们用nvidia-smi观察GPU Util%和htop看CPU负载，找到平衡点。通常num_workers = min(32, 4 * GPU_count)是安全起点。若CPU负载100%而GPU Util<50%，说明worker不足；若CPU负载低而GPU Util仍低，则可能是数据本身IO慢（此时考虑SSD或内存盘）。

3.3 模型构建：支持torch.compile的现代PyTorch写法

PyTorch 2.0的torch.compile是重大升级，但直接套用model = torch.compile(model)可能出问题。Starter Pack的core/model/resnet.py展示了安全集成方式：

class ResNet18(nn.Module): def __init__(self, num_classes: int = 1000, compile_mode: str = "default"): super().__init__() self.backbone = models.resnet18(weights=None) # 不加载预训练权重 self.backbone.fc = nn.Linear(self.backbone.fc.in_features, num_classes) self.compile_mode = compile_mode def forward(self, x: torch.Tensor) -> torch.Tensor: return self.backbone(x) def compile(self, **kwargs) -> "ResNet18": """Safely compile the model with mode-specific defaults.""" if self.compile_mode == "default": # 最安全，默认模式 self.backbone = torch.compile(self.backbone, **kwargs) elif self.compile_mode == "max-autotune": # 极致性能，但编译时间长，适合固定shape self.backbone = torch.compile( self.backbone, mode="max-autotune", fullgraph=True, dynamic=False ) return self

在core/trainer.py中，Trainer类的__init__方法根据配置决定是否编译：

def __init__(self, cfg: dict): # ... 其他初始化 if cfg.get("compile", False): compile_mode = cfg.get("compile_mode", "default") self.model = self.model.compile(compile_mode=compile_mode) logger.info(f"Model compiled with mode: {compile_mode}")

configs/train/default.yaml中控制：

compile: true compile_mode: "default" # 可选: "default", "reduce-overhead", "max-autotune"

注意：torch.compile在训练和推理阶段行为不同。训练时推荐mode="default"（平衡编译时间和性能）；推理时若输入shape固定，可用mode="max-autotune"获得最高吞吐。但切记：dynamic=True（支持变长输入）会显著增加编译时间，且某些算子不支持。

3.4 训练循环：DDP、梯度裁剪、混合精度的无缝集成

core/trainer.py的train_epoch()方法是整个Pack的心脏，它把分布式训练、混合精度、梯度裁剪等复杂逻辑封装成可插拔组件：

def train_epoch(self): self.model.train() total_loss = 0 for batch_idx, (data, target) in enumerate(self.train_loader): data, target = data.to(self.device), target.to(self.device) # 1. 混合精度上下文 with torch.cuda.amp.autocast(enabled=self.use_amp): output = self.model(data) loss = self.criterion(output, target) # 2. 梯度缩放（AMP必需） self.scaler.scale(loss).backward() # 3. 梯度裁剪（防爆炸） if self.cfg.get("grad_clip", 0) > 0: self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_( self.model.parameters(), self.cfg["grad_clip"] ) # 4. 优化器step（含缩放） self.scaler.step(self.optimizer) self.scaler.update() self.optimizer.zero_grad(set_to_none=True) # set_to_none=True节省显存 total_loss += loss.item() return total_loss / len(self.train_loader)

关键细节解析：

• set_to_none=True：比zero_grad()更激进，直接将梯度张量置为None，释放显存。PyTorch 1.9+推荐，实测ResNet18训练显存占用降低12%。

• self.scaler.unscale_(self.optimizer)：在梯度裁剪前，必须先unscale，否则裁剪的是缩放后的梯度（数值巨大，裁剪失效）。

• DDP集成：在cli.py中，我们检测WORLD_SIZE环境变量，自动启用DistributedDataParallel：

  if torch.cuda.device_count() > 1 and int(os.environ.get("WORLD_SIZE", "1")) > 1: model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[local_rank], find_unused_parameters=False # 默认False，除非模型有未参与loss计算的分支 )

find_unused_parameters=False是性能关键。设为True会触发额外的梯度检查，速度下降30%。只有当你确认模型中有分支（如auxiliary head）在某些batch中不参与计算时，才开启。

3.5 日志与检查点：让每一次训练都“可审计”

日志不是为了好看，而是为了“可审计”。Starter Pack的日志系统设计满足三个硬需求：实时可见、历史可查、指标可导出。

实时可见：logging+tqdm深度整合

core/trainer.py中，train_epoch()使用tqdm包裹train_loader，但进度条显示的不仅是batch数，还有实时loss和lr：

pbar = tqdm(self.train_loader, desc=f"Train Epoch {self.epoch}") for batch_idx, (data, target) in enumerate(pbar): # ... 训练逻辑 pbar.set_postfix({ "loss": f"{loss.item():.4f}", "lr": f"{self.optimizer.param_groups[0]['lr']:.6f}" })

同时，logging在INFO级别记录每个epoch的摘要：

logger.info( f"Train Epoch {self.epoch}: " f"Loss={avg_train_loss:.4f} | " f"LR={self.optimizer.param_groups[0]['lr']:.6f} | " f"Time={time.time()-start_time:.2f}s" )

历史可查：结构化日志文件

所有logging输出同时写入logs/train_{timestamp}.log，文件内容严格按时间戳+级别+模块名排序：

2024-05-20 14:22:31,123 INFO [core.trainer] Train Epoch 1: Loss=2.3145 | LR=0.001000 | Time=124.32s 2024-05-20 14:23:05,456 INFO [core.trainer] Val Epoch 1: Acc=72.34% | Best=72.34% | Time=32.11s

这样，grep "Val Epoch" logs/*.log | sort就能得到所有验证结果的时间线。

指标可导出：TensorBoard + CSV双备份

core/trainer.py中，self.writer是SummaryWriter实例，记录所有关键指标：

# 记录scalar self.writer.add_scalar("train/loss", avg_train_loss, self.epoch) self.writer.add_scalar("train/lr", self.optimizer.param_groups[0]['lr'], self.epoch) self.writer.add_scalar("val/acc", val_acc, self.epoch) # 记录histogram（梯度分布，调试用） for name, param in self.model.named_parameters(): if param.grad is not None: self.writer.add_histogram(f"grad/{name}", param.grad, self.epoch)

同时，core/utils.py提供save_metrics_to_csv()函数，将val_acc,val_loss等关键指标追加到metrics.csv：

epoch,train_loss,val_loss,val_acc,best_val_acc 1,2.3145,1.8762,72.34,72.34 2,1.9876,1.7654,75.21,75.21

这个CSV可直接用Excel或Pandas绘图，无需启动tensorboard。

检查点保存：智能覆盖与版本保留

core/trainer.py的save_checkpoint()方法支持两种策略：

• save_best_only: true（默认）：只保存验证指标最佳的模型，文件名为best_model.pt
• save_every_n_epochs: 10：每10个epoch保存一次，文件名为checkpoint_epoch_10.pt

检查点内容包含：

• model_state_dict
• optimizer_state_dict
• scheduler_state_dict
• epoch,best_metric,cfg（配置快照）

注意：cfg被序列化保存，确保未来加载时知道当时的超参。这是复现实验的黄金凭证。

4. 实操全流程：从零启动一个图像分类项目

4.1 环境准备与依赖安装

Starter Pack对环境要求极简，仅需PyTorch 2.0+和基础科学计算库。我们不捆绑conda或docker，因为多数用户已有自己的环境管理习惯。以下是实测通过的安装步骤：

# 创建虚拟环境（推荐） python -m venv pytorch_starter_env source pytorch_starter_env/bin/activate # Linux/Mac # pytorch_starter_env\Scripts\activate # Windows # 安装PyTorch（以CUDA 11.8为例，根据你的GPU选择） pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 安装其他依赖（全部轻量级） pip install hydra-core==1.3.2 omegaconf==2.3.0 tensorboard==2.15.1 tqdm==4.66.1

提示：hydra-core版本锁定在1.3.2，因为1.4+引入了@hydra.main的签名变更，与我们cli.py的@hydra.main(config_path="configs", config_name="train")不兼容。这是经过23次CI失败后确定的稳定组合。

验证安装：

python -c "import torch; print(torch.__version__, torch.cuda.is_available())" # 应输出类似：2.1.0+cu118 True

4.2 数据准备：以Flowers数据集为例

Starter Pack不提供数据下载脚本，因为数据合规性需用户自行确认。我们以公开的Oxford-IIIT Pet Dataset（常被误称为Flowers）为例，说明目录结构和配置编写：

1. 下载并解压数据：

   wget https://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz tar -xzf images.tar.gz # 得到 images/ 目录，内含 37类宠物图片

2. 创建数据配置文件configs/data/pets.yaml：

   _target_: core.data.pets.PetsDataset root: "/path/to/your/images" # 替换为你的实际路径 split: "train" transform: train: - name: "Resize" size: [256, 256] - name: "RandomResizedCrop" size: [224, 224] scale: [0.8, 1.0] - name: "RandomHorizontalFlip" p: 0.5 - name: "ToTensor" - name: "Normalize" mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225] val: - name: "Resize" size: [256, 256] - name: "CenterCrop" size: [224, 224] - name: "ToTensor" - name: "Normalize" mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225]

3. 创建模型配置configs/model/resnet18_pets.yaml：

   _target_: core.model.resnet.ResNet18 num_classes: 37 compile_mode: "default"

4. 创建训练配置configs/train/pets_default.yaml：

   epochs: 50 lr: 0.001 weight_decay: 1e-4 grad_clip: 1.0 compile: true dataloader: batch_size: 64 num_workers: 8 optimizer: _target_: torch.optim.AdamW lr: ${..lr} weight_decay: ${..weight_decay} scheduler: _target_: torch.optim.lr_scheduler.CosineAnnealingLR T_max: ${..epochs}

4.3 启动训练：一条命令，全程可控

一切就绪后，启动训练只需一条命令：

python train.py \ model=resnet18_pets \ data=pets \ train=pets_default \ hydra.run.dir="./outputs/pets_resnet18" \ hydra.job.name="pets_resnet18"

这条命令的含义：

• model=resnet18_pets：加载configs/model/resnet18_pets.yaml
• data=pets：加载configs/data/pets.yaml
• train=pets_default：加载configs/train/pets_default.yaml
• hydra.run.dir：指定输出目录，避免日志混杂
• hydra.job.name：设置job名称，用于tensorboard tag

训练过程中，你会看到：

• 实时tqdm进度条，显示loss和lr
• 控制台INFO日志，记录epoch摘要
• outputs/pets_resnet18/目录下生成：
  • train_2024-05-20_14-22-31.log：结构化日志文件
  • metrics.csv：指标CSV
  • events.out.tfevents.*：tensorboard事件文件
  • best_model.pt：最佳模型检查点

4.4 验证与推理：如何用训练好的模型做预测

Starter Pack提供infer.py脚本，用于单张图片或批量推理。以验证best_model.pt为例：

python infer.py \ --model_path outputs/pets_resnet18/best_model.pt \ --config_path outputs/pets_resnet18/.hydra/config.yaml \ --image_path path/to/test_image.jpg \ --top_k 3

infer.py的核心逻辑：

1. 加载config.yaml，重建完全相同的transform和model结构
2. 加载best_model.pt的state_dict，严格匹配键名（strict=True）
3. 对输入图片应用val阶段的transform
4. 模型前向，输出top-k预测类别和置信度

实操心得：infer.py不依赖训练时的core/目录，它通过config.yaml中的_target_字段动态导入类，因此即使你移动了代码位置，只要配置正确，推理依然有效。这是“配置即契约”的体现。

5. 常见问题与避坑指南：那些文档里不会写的血泪教训

5.1 “CUDA out of memory” —— 显存爆炸的10种可能与定位法

OOM是PyTorch新手第一道墙。Starter Pack内置了三重防御，但你仍需知道如何破局：

防御一：torch.cuda.memory_summary()

在core/trainer.py的train_epoch()开头，我们添加了内存快照：

if self.epoch == 1 and batch_idx == 0: logger.info(torch.cuda.memory_summary())

这会在第一个batch前打印显存分配详情，包括：

• allocated：当前分配的显存（GB）
• reserved：CUDA driver预留的显存（通常>allocated）
• active：活跃块数量
• inactive：可回收但未释放的块

防御二：torch.cuda.max_memory_allocated()

在validate_epoch()后，记录峰值显存：

peak_mem = torch.cuda.max_memory_allocated() / 1024**3 logger.info(f"Peak GPU memory: {peak_mem:.2f} GB")

防御三：torch.utils.checkpoint（梯度检查点）

当模型太大时，在core/model/resnet.py中启用检查点：

from torch.utils.checkpoint import checkpoint class ResNet18(nn.Module): def forward(self, x): x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x = self.maxpool(x) # 对每个layer应用checkpoint x = checkpoint(self.layer1, x) x = checkpoint(self.layer2, x) x = checkpoint(self.layer3, x) x = checkpoint(self.layer4, x) x = self.avgpool(x) x = torch.flatten(x, 1) x = self.fc(x) return x

注意：checkpoint会增加约15%计算时间，但显存占用可降50%。仅在allocated接近reserved时启用。

OOM定位速查表