Проблема несовместимости формы с сеткой плотности смеси в tenorflow/keras

Question

Я застрял в проблеме при вычислении сети плотности смеси через keras (с использованием бэкэндов тензорного потока). Цель этого MDN - изучить скрытое представление изображения (чтобы реализовать предсказания MDN в автоэнкодере). Затем я хотел бы смоделировать свое входное изображение как многомерное нормальное распределение и получить в качестве сетевого выхода вектор mu и sigma (размером 64 каждый) и набор из N весов альфа (где N - количество компонентов в смеси.). Если я рассматриваю выход, имеющий форму 64 для каждого параметра, все работает хорошо, но нет смысла получать больше альфа-факторов, чем компонентов (в моем случае это альфа-размер 64). Когда я пытаюсь указать альфа, чтобы получить форму, отличную от mu и sigma, возникают некоторые проблемы.

Чтобы получить ковариационную матрицу для реализации в модуле MixtureSameFamily тензорного потока, я рассматриваю диагональную матрицу сигма-вектора. Затем я нашел эту функцию потерь (с отрицательной логической вероятностью) на многих форумах и попытался адаптировать ее к моей проблеме:

def slice_parameter_vectors(parameter_vector):
    return tf.split(parameter_vector,[1*components,64*components,64*components],axis=1)

def gnll_loss(y, parameter_vector):
    alpha, mu, sigma = slice_parameter_vectors(parameter_vector)  # Unpack parameter vectors

gm = tfd.MixtureSameFamily(
    mixture_distribution=tfd.Categorical(probs=alpha),
    components_distribution=tfd.MultivariateNormalDiag(
        loc=mu,      
        scale_diag=  sigma))

    log_likelihood = gm.log_prob(tf.transpose(y))                 # Evaluate log-probability of y

    return -tf.reduce_mean(log_likelihood, axis=-1)

Если я попытаюсь наполнить свою сеть некоторыми данными и скомпилировать их, у меня всегда будет эта ошибка:

InvalidArgumentError: Incompatible shapes: [64,1,2] vs. [2,64]
     [[{{node loss_20/concatenate_6_loss/MultivariateNormalDiag/log_prob/affine_linear_operator/inverse/sub}} = Sub[T=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:GPU:0"](loss_20/concatenate_6_loss/MixtureSameFamily/log_prob/pad_sample_dims/Reshape, loss_20/concatenate_6_loss/split:1)]]
     [[{{node loss_20/mul/_4221}} = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_1846_loss_20/mul", tensor_type=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

64 - размерность вектора, а 2 - размер пакета. Использование размера пакета 1 работает, но дает мне некоторую потерю NaN.

Вот как создаются слои альфа, му и сигма:

    fc    = Dense((no_parameters-1) * components*64 + components, activation="tanh", name="fc")(layer)
    alphas = Dense(1*components, activation="softmax", name="alphas")(fc)
    mus    = Dense(64*components, name="mus")(fc)
    sigmas = Dense(64*components, activation=nnelu, name="sigmas")(fc)
    pvec   = Concatenate(axis=1)([alphas,mus,sigmas])
    mdn    = Model(inputs=inputs,outputs=pvec)

Таким образом, возникает вопрос: можно ли это сделать с помощью тензорного потока? Есть ли здесь кто-нибудь еще, использующий такие сети, и мог бы объяснить мне, как бороться с функцией потерь?

С уважением,

Адриан