В настоящее время я воссоздаю AlphaZero Deepmind в качестве личного проекта с использованием Keras с бэкэндом TensorFlow. Во время каждой симуляции игры NN, который я использую, должен делать много прогнозов. NN работает должным образом, но использует чрезмерное количество оперативной памяти.
Я использую psutil.Process(os.getpid()).memory_info().rss для печати использования памяти, а когда model.predict() вызывается впервые, использование памяти увеличивается с 0,3 до 1,6 ГБ. При последовательных вызовах новая память (или, по крайней мере, очень маленькая) не выделяется.
Я не уверен, что это ожидаемое поведение, но я бы хотел, чтобы программа работала настолько эффективно, насколько это возможно, так как я хотел бы запускать ее на TPU в Google Compute Engine.
Ниже приведен код для моей модели 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}
)
Нейронная сеть очень велика, поэтому ее размер может влиять на объем выделяемой оперативной памяти, но 1,6 ГБ кажется неоправданным. Кроме того, если я уменьшу модель до наличия только 2 остаточных слоев вместо 19, она все равно выделит около 1,54 ГБ вместо 1,6 ГБ, использовавшихся ранее.