python - 为什么单次调用 Keras model.predict() 使用 1.3 GB 内存?
问题描述
我目前正在使用 Keras 和 TensorFlow 后端将 Deepmind 的 AlphaZero 重新创建为个人项目。在每次游戏模拟过程中,我使用的 NN 都需要进行许多预测。NN 正常工作,但使用了过多的 CPU RAM。
我psutil.Process(os.getpid()).memory_info().rss
用来打印内存使用情况,model.predict()
第一次调用时,内存使用量从大约 0.3 GB 跳到 1.6 GB。在连续调用时,不会分配新的内存(或至少很少)。
我不确定这是否是预期行为,但我希望程序尽可能高效地运行,就像我希望在 Google Compute Engine 上的 TPU 上运行它一样。
下面是我的 Keras 模型的代码:
main_input = Input(shape = self.input_dim, name = 'main_input', dtype="float32")
x = self.conv_layer(main_input, 256, (3,3))
for i in range(19): # AlphaZero used 19
x = self.residual_layer(x, 256, (3, 3))
vh = self.value_head(x)
ph = self.policy_head(x)
model = Model(inputs=[main_input], outputs=[vh, ph])
model.compile(loss={'value_head': 'mean_squared_error', 'policy_head': softmax_cross_entropy_with_logits},
optimizer=SGD(lr=self.learning_rate, momentum = MOMENTUM), #config.MOMENTUM
loss_weights={'value_head': 0.5, 'policy_head': 0.5}
)
神经网络非常大,因此它的大小可能会影响分配的 RAM 量,但 1.6 GB 似乎不合理。此外,如果我将模型减少到只有 2 个剩余层而不是 19 个,它仍然会分配大约 1.54 GB 而不是之前使用的 1.6 GB。
解决方案
推荐阅读
- git-push - Git 强制推动 CI 构建挂起
- spring-boot - 运行spring boot测试时出错,找不到SpringRunner.class
- python-3.6 - 使用 Python3 对重复项进行排序和删除
- eclipse - Eclipse 无法从主机操作系统识别 m2 存储库
- cuda - NVCC unistd.h(792):错误:需要一个标识符
- ios - 如何在 Swift 中隐藏 UITableView?
- object - Access VBA 如何在记录集中查找值,此类对象不支持错误代码“3251”操作
- r - 从 .Rnw 到 .Rmd 的自动转换
- javascript - 加载组件时出现奇怪的错误
- javascript - 在新窗口中打开聊天小部件