PyTorch中自定义LSTM模型保存后体积过大的原因咨询

1. 非参数的隐藏状态张量被保存了 #

你的模型里定义了self.hidden和self.hiddenc，这俩是LSTM的隐藏状态，不属于可训练参数（所以model.named_parameters()扫不到它们），但用torch.save(model)保存整个模型时，会把模型实例的所有属性都打包进去，包括这两个张量。

要是init_hidden()返回的是(num_layers, batch_size, hidden_size)形状的张量，比如训练时用的batch_size是256，单个隐藏状态就有3*256*512=393216个元素，每个float32类型占4字节，单个张量就约1.5MB，两个加起来3MB；如果batch_size更大（比如1024），这俩张量的体积能超过12MB，再加上模型参数本身的大小，总容量很容易达到30MB。

2. 模型参数实际体积远超你的预期 #

你以为三层LSTM只有几十KB，其实参数数量得这么算：

• 每个LSTM层对应4组权重/偏置（输入门、遗忘门、细胞更新、输出门），单LSTM层的参数数是 4 * hidden_size * (input_size + hidden_size)
• 三层LSTM总参数就是 3 * 4 * 512 * (obs_size + 512)
• 再加上线性层i2o的 512*256 + 256 个参数

假设你的obs_size是128，仅LSTM部分就有3*4*512*(128+512)=3932160个float32元素，体积约15.7MB，加上线性层的512KB，参数本身就接近16MB，再加上隐藏状态的体积，总大小轻松达到30MB。

解决办法 #

• 不要将隐藏状态设为模型属性：训练时临时生成隐藏状态即可，不用存在模型实例中，这样保存时就不会带上这些冗余张量
• 保存模型状态字典而非整个模型：使用torch.save(model.state_dict(), "model.pth")，只会保存可训练参数和必要的缓冲区（比如BN层的均值），不会存储模型的其他无关属性，体积会大幅缩小

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