使用torch.autograd.profiler.profile时发生类型错误

Question

我正试图分析我的模型的内存消耗，如下所述：https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html

使用这些线路：

with profiler.profile(profile_memory=True, record_shapes=True) as prof:
tubes, _, _ = zip(*model(imgs, img_metas, return_loss=False))
print(prof.key_averages().table(sort_by="self_cpu_memory_usage", row_limit=10))

但是我收到这个错误

TypeError:init((得到了一个意外的关键字参数"profile_memory">

那么，这个错误的原因和/或解决方案是什么？

Answer 1

您需要升级您的PyTorch版本。查看代码，可以看到profile_memory参数是在PyTorch v1.6.0中首先添加到函数签名中的。

您也可以通过torch.autograd的文档看到这一点。您可以看到PyTorch v1.1.0中没有该参数。