Я работаю над вариационным авто-кодировщиком, и я бы хотел, чтобы ранее использовавшиеся в регуляризации KL-расхождения скрытого распределения обновили lo c (среднее) и scale (stddev).
Приведенный ниже фрагмент - это надуманный минимальный пример, демонстрирующий то, чего я пытаюсь достичь. Это начинает работать, но затем просто замирает после некоторого случайного числа эпох (иногда 1, иногда 200, но обычно около 7 или 8). Там нет ни сообщения об ошибке, ни чего-либо.
loc = tf.Variable(tf.random.normal([ndim], stddev=0.1, dtype=tf.float32))
scale = tfp.util.TransformedVariable(
tf.random.normal([ndim], mean=1.0, stddev=0.1, dtype=tf.float32),
bijector=tfb.Chain([tfb.Shift(1e-5), tfb.Softplus(), tfb.Shift(0.5413)]))
prior = tfd.Independent(tfd.Normal(loc=loc, scale=scale), reinterpreted_batch_ndims=1)
_input = tfkl.Input(shape=(1,))
_loc = tfkl.Dense(ndim, name="loc_params")(_input)
_scale = tfkl.Dense(ndim, name="untransformed_scale_params")(_input)
_scale = tf.math.softplus(_scale + np.log(np.exp(1) - 1)) + 1e-5
_output = tfpl.DistributionLambda(
make_distribution_fn=lambda t: tfd.Independent(tfd.Normal(loc=t[0], scale=t[1])),
activity_regularizer=tfpl.KLDivergenceRegularizer(prior, use_exact_kl=True, weight=0.1)
)([_loc, _scale])
model = tf.keras.Model(_input, _output)
model.compile(optimizer='adam', loss=lambda y_true, model_out: -model_out.log_prob(y_true))
hist = model.fit(ds, epochs=N_EPOCHS, verbose=2)
У меня есть работающий гист здесь .
Более конкретный пример и архитектура, близкая к тому, что я пытаюсь обновить и упростить, это пример tfp для disentangled_vae. В его ручном обучении l oop новый tfd.MultivariateNormalDiag создается для каждого l oop, хотя он параметризован с использованием постоянных tf.Variables. Я стараюсь изо всех сил избегать циклов ручного обучения, и я также пытаюсь перейти к более похожему на Keras синтаксису, поэтому я бы предпочел не делать прямой порт этого примера.
Любой совет с благодарностью Спасибо!
Редактировать: Activity_regularizer, кажется, работает нормально при подключении к латентному (узкому) распределению. У меня есть более полный пример в этой записной книжке Colab . Поскольку это работает в моей архитектуре, я больше не нуждаюсь в ответе.
Однако я очень сомневаюсь, что использование модели, подходящей для замораживания, является желательным поведением, поэтому это остается проблемой.