PyTorch训练绘图时CUDA tensor转numpy报错修复方法

报错是因为你记录训练/验证损失、准确率时，保存的是驻留在CUDA显存中的PyTorch张量。matplotlib绘图依赖numpy数组，而CUDA设备上的张量无法直接转换为numpy格式，必须先将张量拷贝到主机CPU内存才能完成格式转换。

只需要在存储绘图用的指标数据、转换numpy格式前，调用.cpu()方法将对应张量迁移到CPU即可。注意不要在训练前向、反向传播环节随意调用.cpu()，会打断GPU计算流拖慢训练速度，仅在每个epoch结束统计完整体指标后做迁移即可。

错误写法参考（会触发该报错） #

# 每个epoch统计完的指标仍是CUDA张量
train_loss = total_train_loss / len(train_loader)
train_acc = total_train_correct / len(train_loader.dataset)
val_loss = total_val_loss / len(val_loader)
val_acc = total_val_correct / len(val_loader.dataset)

# 直接对CUDA张量调用.numpy()会触发报错
train_loss_curve.append(train_loss.numpy())
train_acc_curve.append(train_acc.numpy())
val_loss_curve.append(val_loss.numpy())
val_acc_curve.append(val_acc.numpy())

正确修改写法 #

# 每个epoch统计完的指标仍是CUDA张量
train_loss = total_train_loss / len(train_loader)
train_acc = total_train_correct / len(train_loader.dataset)
val_loss = total_val_loss / len(val_loader)
val_acc = total_val_correct / len(val_loader.dataset)

# 先调用.cpu()迁移到主机内存，再转numpy存入绘图用的列表
train_loss_curve.append(train_loss.cpu().numpy())
train_acc_curve.append(train_acc.cpu().numpy())
val_loss_curve.append(val_loss.cpu().numpy())
val_acc_curve.append(val_acc.cpu().numpy())

更省显存的写法 #

如果统计指标时直接用.item()取张量的标量值，该方法会自动把CUDA张量上的数值拷贝到CPU返回Python原生浮点数/整数，后续不需要额外转numpy也不会触发报错：

# 逐batch累加时直接取原生数值
running_train_loss += loss.item()
running_train_correct += (preds == labels).sum().item()

提示：如果你之前把每个batch的损失、正确数都以CUDA张量的形式存在列表里，记得统计完epoch结果后清空这部分无用缓存，避免显存占用持续升高。

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