Блог

В настоящее время я воссоздаю AlphaZero Deepmind в качестве личного проекта, используя Keras с серверной частью TensorFlow. Во время каждого игрового моделирования NN, который я использую, должен делать много прогнозов. NN работает так, как должен, но использует чрезмерный объем оперативной памяти процессора.

Я использую psutil.Process(os.getpid()).memory_info().rss для печати использования памяти, и когда model.predict() вызывается в первый раз, использование памяти увеличивается примерно с 0,3 ГБ до 1,6 ГБ. При последовательных вызовах новая память (или, по крайней мере, очень мало) не выделяется.

Я не уверен, что это ожидаемое поведение, но я хотел бы, чтобы программа работала как можно эффективнее, поскольку я хотел бы запустить ее на процессорных процессорах в 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 ГБ, использовавшихся ранее.