Я пытаюсь реорганизовать некоторые из наших методов обучения с подкреплением, используя tf.keras. Модели keras, кажется, довольно хорошо работают для сетей с прямой связью, которые я использую для функций-значений и препроцессоров, но я пытаюсь реализовать некоторые вероятностные модели (например, политики) с помощью keras. В частности, объединение tfp.distributions и tfp.bijectors с tf.keras.Model кажется мне очень не интуитивным, и все примеры, которые я видел (например, [1, 2, 3]), либо слишком упрощены, либо полагаются на «хаки», которые, на мой взгляд, подрывают многие преимущества использования моделей keras (например, возможность скрыть обработку ввода, сеансы и числовую оценку от самой модели).
Предположим, я хочу реализовать политику скрытого пространства, как описано в [4], которая использует поток RealNVP для преобразования гауссовых выборок, обусловленных состояниями, в действия. Политика должна поддерживать как минимум две операции:
1. Выборка действий Y такая, что
Y = g (X | S)
X ~ Normal (0, 1),
где g - преобразование RealNVP, описанное в [4, 5], а S - переменные состояния (например, наблюдения за состоянием в случае RL).
2. Вычисление логарифмических вероятностей выбранных Y.
Простая реализация может выглядеть примерно так:
class LearnableConditionalRealNVP(object):
def __init__(self, input_shape, output_shape):
self._input_shape = input_shape
self._output_size = np.prod(output_shape)
conditions = tf.keras.layers.Input(shape=input_shape)
batch_size = tf.keras.layers.Lambda(
lambda x: tf.shape(x)[0])(conditions)
def samples_and_log_probs_fn(inputs):
conditions, batch_size = inputs
base_distribution = tfp.distributions.MultivariateNormalDiag(
loc=tf.zeros(output_shape),
scale_diag=tf.ones(output_shape))
real_nvp_bijector = tfp.bijectors.RealNVP(
num_masked=self._output_size // 2,
shift_and_log_scale_fn=conditioned_real_nvp_template(
hidden_layers=(128, 128),
activation=tf.nn.relu),
name='real_nvp')
distribution = (
tfp.distributions.ConditionalTransformedDistribution(
distribution=base_distribution,
bijector=real_nvp_bijector))
samples = distribution.sample(batch_size)
log_probs = distribution.log_prob(samples)
return [samples, tf.reshape(log_probs, (-1, 1))]
samples, log_probs = tf.keras.layers.Lambda(
samples_and_log_probs_fn)([conditions, batch_size])
self.samples_and_log_probs_model = tf.keras.Model(
conditions, [samples, log_probs])
def samples_and_log_probs(self, conditions):
return self.samples_and_log_probs_model(conditions)
def samples_and_log_probs_np(self, conditions):
return self.samples_and_log_probs_model.predict(conditions)
где conditioned_real_nvp_template создает сеть с прямой связью, которая объединяет скрытые выборки и значения условий вдоль последней оси и использует их в качестве входных данных. Полный пример можно найти здесь: https://gist.github.com/hartikainen/17ac2ec102032e986cb4d31e225f592a
Этот способ обработки дистрибутива принес бы мне два основных преимущества. Во-первых, мне не нужно обрабатывать повторное использование параметров вручную. Я могу вызывать samples_and_log_probs в своем коде несколько раз, и он автоматически использует параметры модели. Во-вторых, мне не нужно ничего знать о сессиях, если я хочу получить числовые результаты. Заключение промежуточных слоев в их собственные модели позволяет мне вызывать методы прогнозирования, которые обрабатывают сеансы.
Проблемы возникают при расширении примера еще немного. Предположим, что я хотел изменить LearnableConditionalRealNVP так, чтобы я мог предоставить скрытые выборки x в качестве входных данных для него, и вместо вызова distribution.sample() в samples_and_log_probs_fn он вернул бы distribution.forward(x). Или, может быть, я хочу получить образцы и log_probs отдельно от модели. Это потребовало бы, чтобы я разделил samples_and_log_probs_fn на две отдельные лямбда-функции, но сделать это не тривиально, если я хочу поделиться параметрами бижектора RealNVP (потому что я не могу передать биектор как вход / выход в / из слой керас).
Я пытался решить эти проблемы путем подкласса LearnableConditionalRealNVP из tf.keras.Model, но все мои попытки приводили к грязным реализациям, в основном из-за изменения входов и выходов. В частности, я не смог создать call -метод для модели так, чтобы модель сохранила возможность использования с predict, и мне пришлось бы делать некоторые трюки с __call__ -методом. Ничто из этого не кажется ужасным, но они все же увеличивают накладные расходы на использование моделей keras настолько, что мне легче реализовать эти типы вещей в простом тензорном потоке и обрабатывать сеансы, пустые выводы и т. Д. Вручную.
Мои вопросы:
- Предполагается, что биекторы / распределения тензорного потока совместимы с кератами (моделями)? Если да, кто-нибудь знает, есть ли какие-нибудь нетривиальные примеры, на которые я мог бы взглянуть? Если нет, планируется ли сделать их совместимыми?
- Как мне использовать модели keras в тех случаях, когда то, что я считаю одной моделью, имеет несколько разных выходов? Например, в моем примере выше, обучаемый дистрибутив RealNVP интуитивно чувствует, что он должен быть единой моделью, но в то же время он имеет несколько, возможно, независимых входов / выходов, что затрудняет встраивание в каркас модели keras. Я готов признать, что моя интуиция здесь неправильна, и в этом случае было бы неплохо услышать, что является лучшей практикой для построения таких моделей.
- Существует ли способ передачи нетензорных данных в качестве входных / выходных данных в / из модели keras, как это сделано в [1], при сохранении каким-либо образом связанных моделей. Если вместо
tfe.Variable s вы используете входные данные в этом примере, он ломается, потому что график не связан.
Редактировать: после публикации этого и выполнения дополнительного тестирования с вышеупомянутой реализацией, я заметил, что эта модель, в конце концов, не обучаема, так как переменные для биектора RealNVP создаются в слое keras lambda. Это наводит меня на мысль о том, что функциональный способ построения этих моделей вообще нельзя использовать для моделей такого типа.
[1] https://github.com/tensorflow/probability/blob/5f5510201865350b6cce2a0f18fbe0cdf4f15eee/tensorflow_probability/examples/disentangled_vae.py#L186
[2] https://blog.keras.io/building-autoencoders-in-keras.html
[3] http://louistiao.me/posts/implementing-variational-autoencoders-in-keras-beyond-the-quickstart-tutorial/
[4] https://arxiv.org/pdf/1804.02808.pdf
[5] https://arxiv.org/abs/1605.08803