Google Colab运行mobile-yolov5-pruning-distillation项目遇RuntimeError求助

错误原因 #

报错核心是：需要梯度的叶子张量（模型可训练参数）不允许执行原地修改操作。这里的b是YOLO模型输出层的偏置参数，属于叶子张量，+=是原地修改操作，破坏了PyTorch的梯度计算图追踪机制，因此触发RuntimeError。

两种可行的修改方案 #

方案1：使用torch.no_grad()包裹初始化代码 #

直接在原地操作的代码外层加上torch.no_grad()上下文管理器，临时停止梯度追踪，允许对叶子张量做原地修改：

with torch.no_grad():
    b[:, 4] += math.log(8 / (640 / s) ** 2)  # obj (8 objects per 640 image)

这种方法最简单，因为模型偏置初始化本来就不需要计算梯度，不会影响后续训练的梯度回传。

方案2：创建副本修改后赋值回原张量 #

如果不想依赖上下文管理器，可以先克隆张量副本，修改后再赋值回原参数：

# 克隆张量副本，脱离原计算图追踪
b_modified = b.clone()
b_modified[:, 4] += math.log(8 / (640 / s) ** 2)
# 在无梯度上下文里将修改后的值复制回原参数
with torch.no_grad():
    b.copy_(b_modified)

验证修改效果 #

修改后重新运行train.py，初始化模型时就不会再触发该RuntimeError，后续训练流程可以正常执行。

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