运行LightXML脚本遇CUDA内存不足，无法找到batch size配置

问题概述 #

运行LightXML脚本时触发CUDA内存不足错误：

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 192.00 MiB (GPU 0; 6.00 GiB total capacity; 4.71 GiB already allocated; 0 bytes free; 4.82 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

已知需减小batch size，但找不到其定义位置，脚本卡在train_loss = model.one_epoch(...)调用中，具体停在one_epoch方法的outputs = self(**inputs)处，最终由forward函数的outs=self.bert()步骤触发内存溢出。

定位Batch Size的定义 #

LightXML的batch size在数据加载器（DataLoader）初始化代码中设置，也就是创建trainloader、validloader、testloader的位置，按以下步骤查找：

1. 找到脚本中主函数（如main.py）或数据处理模块（如data_utils.py）里创建DataLoader的代码段
2. 定位类似如下的代码：

   trainloader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
   validloader = DataLoader(valid_dataset, batch_size=32, num_workers=4)
   

   这里的batch_size参数就是需要调整的数值，尝试将其调小（比如从32改为16、8）

额外内存优化方案 #

若调整batch size后仍内存不足，可尝试以下操作：

• 更换轻量化预训练模型：比如用bert-small替代bert-base，减少模型参数占用
• 调整梯度累积：代码中已有loss /= self.update_count的逻辑，增大update_count的值可间接降低有效batch size
• 清理GPU缓存：在check_memory()函数中添加torch.cuda.empty_cache()语句，释放闲置内存
• 禁用混合精度训练：若代码中启用了AMP，可尝试注释掉with amp.scale_loss(loss, optimizer) as scaled_loss:相关逻辑

内容的提问来源于stack exchange，提问作者Katia