### 解决 PyTorch 模型训练中 GPU 内存不足的方法和优化技巧 #### 1. 减少批量大小 (Batch Size) 减少输入数据的批量大小可以显著降低每次迭代所需的显存。虽然这可能会影响模型收敛速度，但在资源有限的情况下是一个有效的解决方案[^1]。 ```python batch_size = 32 # 原始批大小 new_batch_size = batch_size // 2 # 尝试减半 ``` #### 2. 使用混合精度训练 (Mixed Precision Training) 通过采用 FP16 数据类型代替默认的 FP32 可节省一半以上的显存空间，并且通常不会影响最终模型性能。PyTorch 提供了 `torch.cuda.amp` 工具来简化这一过程。 ```python from torch.cuda import amp scaler = amp.GradScaler() for input, target in data_loader: optimizer.zero_grad() with amp.autocast(): output = model(input) loss = criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() ``` #### 3. 清理不必要的张量对象 及时释放不再使用的变量有助于回收宝贵的 GPU 资源。可以通过调用 `.detach()` 方法断开计算图连接并手动删除不需要的对象。 ```python del intermediate_variable import gc; gc.collect() torch.cuda.empty_cache() ``` #### 4. 利用分布式训练或多卡设置 如果有多块 GPU 可用，则可考虑利用多进程或多线程的方式分配工作负载到不同的设备上执行，从而分摊单个 GPU 的压力。 ```python import torch.distributed as dist dist.init_process_group(backend='nccl') model = nn.parallel.DistributedDataParallel(model) ``` #### 5. 合理设计网络结构 避免过度复杂的层堆叠或参数过多的设计方案，在满足业务需求的前提下尽可能精简模型规模以适应硬件条件限制。 ---