Как я могу заставить настраиваемую функцию потерь работать с многопроцессорностью?

Question

Я пытаюсь создать простую версию AlphaZero для крестиков-ноликов с помощью Keras. Это работает, но создание данных для самостоятельного воспроизведения происходит очень медленно, поэтому я решил распараллелить его, используя многопроцессорность, чтобы создать пакет обучающих данных в одном go.

Однако я использую настраиваемую потерю функция, которую он нормально обрабатывает, но с многопроцессорной обработкой она дает мне эту ошибку:

Traceback (most recent call last):
  File "C:\Users\hughs\AppData\Local\Programs\Python\Python37\lib\multiprocessing\process.py", line 297, in _bootstrap
    self.run()
  File "C:\Users\hughs\AppData\Local\Programs\Python\Python37\lib\multiprocessing\process.py", line 99, in run
    self._target(*self._args, **self._kwargs)
  File "C:\Users\hughs\AppData\Local\Programs\Python\Python37\lib\multiprocessing\pool.py", line 110, in worker
    task = get()
  File "C:\Users\hughs\AppData\Local\Programs\Python\Python37\lib\multiprocessing\queues.py", line 354, in get
    return _ForkingPickler.loads(res)
  File "C:\Users\hughs\PycharmProjects\Chess\venv\lib\site-packages\keras\engine\network.py", line 1334, in __setstate__
    model = saving.unpickle_model(state)
  File "C:\Users\hughs\PycharmProjects\Chess\venv\lib\site-packages\keras\engine\saving.py", line 604, in unpickle_model
    return _deserialize_model(h5dict)
  File "C:\Users\hughs\PycharmProjects\Chess\venv\lib\site-packages\keras\engine\saving.py", line 369, in _deserialize_model
    sample_weight_mode=sample_weight_mode)
  File "C:\Users\hughs\PycharmProjects\Chess\venv\lib\site-packages\keras\backend\tensorflow_backend.py", line 75, in symbolic_fn_wrapper
    return func(*args, **kwargs)
  File "C:\Users\hughs\PycharmProjects\Chess\venv\lib\site-packages\keras\engine\training.py", line 119, in compile
    self.loss, self.output_names)
  File "C:\Users\hughs\PycharmProjects\Chess\venv\lib\site-packages\keras\engine\training_utils.py", line 828, in prepare_loss_functions
    loss_functions = [get_loss_function(l) for l in loss]
  File "C:\Users\hughs\PycharmProjects\Chess\venv\lib\site-packages\keras\engine\training_utils.py", line 828, in <listcomp>
    loss_functions = [get_loss_function(l) for l in loss]
  File "C:\Users\hughs\PycharmProjects\Chess\venv\lib\site-packages\keras\engine\training_utils.py", line 705, in get_loss_function
    loss_fn = losses.get(loss)
  File "C:\Users\hughs\PycharmProjects\Chess\venv\lib\site-packages\keras\losses.py", line 795, in get
    return deserialize(identifier)
  File "C:\Users\hughs\PycharmProjects\Chess\venv\lib\site-packages\keras\losses.py", line 776, in deserialize
    printable_module_name='loss function')
  File "C:\Users\hughs\PycharmProjects\Chess\venv\lib\site-packages\keras\utils\generic_utils.py", line 167, in deserialize_keras_object
    ':' + function_name)
ValueError: Unknown loss function:softmax_cross_entropy_with_logits_v2

Вот упрощенная версия моего кода:

from multiprocessing import Pool
import numpy


def play_game(neural_net):
    foo = [[0 for i in range(3)] for j in range(3)]
    return neural_net.predict(numpy.array([foo]))


def make_net():
    from keras import Input, Model
    from keras.layers import Dense
    from keras.optimizers import Adam
    from tensorflow import nn

    input = Input((3, 3))
    output = Dense(1)(input)
    model = Model(inputs=input, outputs=[output])
    adam = Adam(0.2)
    model.compile(adam, [nn.softmax_cross_entropy_with_logits])
    model.summary()
    return model


def main():
    neural_net = make_net()
    with Pool() as p:
        training_data = p.map(play_game, [neural_net for i in range(10)], 3)
    print(training_data)


if __name__ == '__main__':
    main()

Я пробовал это с разными функции потерь являются внутренними для keras, и это решает проблему (но дает мне другую ошибку, чем что-то еще). Я также попытался включить custom_objects в вызываемые функции нейронной сети, но это не дало результата. Любая помощь будет принята с благодарностью.