加载操作后 Pytorch GPU 内存增加答案

【问题标题】：Pytorch GPU memory increase after load operation加载操作后 Pytorch GPU 内存增加
【发布时间】：2019-09-23 23:38:12
【问题描述】：

我有一个大小为 386MB 的 pytorch 模型，但是当我加载模型时

state = torch.load(f, flair.device)

我的 GPU 内存占用高达 900MB，为什么会出现这种情况？有没有办法解决这个问题？

这就是我保存模型的方式

model_state = self._get_state_dict()

# additional fields for model checkpointing
model_state["optimizer_state_dict"] = optimizer_state
model_state["scheduler_state_dict"] = scheduler_state
model_state["epoch"] = epoch
model_state["loss"] = loss

torch.save(model_state, str(model_file), pickle_protocol=4)

【问题讨论】：

标签： pytorch

【解决方案1】：

可能是optimizer_state 占用了额外的空间。一些优化器（例如 Adam）跟踪每个可训练参数的统计信息，例如一阶和二阶矩。如您所知，此信息占用空间。

你可以先加载到CPU：

state = torch.load(f, map_location=torch.device('cpu'))

【讨论】：

推理的时候不需要加载吧？我的意思是有没有办法我不能加载优化器状态？或者我现在坚持这个？
@Ryan 您可以加载到 CPU 而不是直接加载到 GPU，然后仅将模型移动到 GPU。
最后一个问题，我该怎么做？