InvalidArgumentError: пользовательский слой несовместимых фигур и пользовательские кера потери журнала

Question

Я пытаюсь создать рекуррентную сеть плотности, аналогичную Grave, для прогнозирования почерка: https://arxiv.org/abs/1308.0850

Я хочу использовать только керас и вероятность тензорного потока для выполнения работы. Тем не менее, я получаю следующую ошибку InvalidArgumentError: Incompatible shapes: [6,20,2,1] vs. [6,20,24,2] [[{{node MultivariateNormalDiag_11/log_prob/affine_linear_operator/inverse/sub}} = _MklSub[T=DT_FLOAT, _kernel="MklOp", _device="/job:localhost/replica:0/task:0/device:CPU:0"](MixtureSameFamily_16/log_prob/pad_sample_dims/Reshape, reshape_33/Reshape, MixtureSameFamily_16/log_prob/pad_sample_dims/Reshape:1, reshape_33/Reshape:1)]]

Мой пример вывода должен иметь размерность:

[6,20,2] за [batch_size, sequence_length, input_dim] Кажется, это была подушка.

[6,20,24,2] относится к [batch_size, sequence_length, number_mixture, input_dim]

[batch_size, sequence_length, input_dim]

Вот подход, я начинаю определять нейронную сеть для вычисления альфа, mus и sigma:

def neural_net(x):
  # Getting rnn input shape
  seq_len, input_dim = x.shape[1:]

  hidden = GRU(hidden_units,input_shape=(seq_len ,input_dim),return_sequences=True)(x)
  alpha = Dense(k_mixt, activation='softmax',name='alpha')(hidden)
  sigma = Dense(k_mixt, activation = 'exponential',name='sigma')(hidden)
  mu = Dense(2*k_mixt,name='mu')(hidden)
  return alpha, mu, sigma

После этого я делаю некоторые изменения. В некотором смысле, потому что мы хотим использовать двумерный гауссов, и у нас есть один параметр для каждой ковариационной матрицы.

def process_outputs(mu,sigma):
  # Reshaping the outputs of the neural net
  mu = Reshape([seq_len,k_mixt,input_dim])(mu)
  sigma_reshaped = Reshape([seq_len,k_mixt,1])(sigma) 
  sigma_out = Concatenate()([sigma_reshaped,sigma_reshaped])
  return mu, sigma_out

Чтобы правильно использовать многомерную нормаль вероятности тензорного потока, нам нужно иметь одинаковое количество аргументов для mu и sigma.

def gaussian_mixture(alpha,mu,sigma):
  gaussians = tfd.MultivariateNormalDiag(loc = mu, scale_diag = sigma)
  alpha_prob = tfd.Categorical(probs = alpha)

  gm = tfd.MixtureSameFamily(
            mixture_distribution=alpha_prob,
                    components_distribution=gaussians)  
  return gm

Если мы хотим использовать структуру te keras, тем не менее, нам понадобится слой keras для создания гауссовой смеси, иначе мы не сможем пропустить градиенты, поэтому я определю пользовательский слой.

class GaussianMixture(Layer):

  def __init__(self, **kwargs):
    super(GaussianMixture, self).__init__(**kwargs)

  def build(self, input_shape):
    super(GaussianMixture, self).build(input_shape)  

  def call(self, inputs):
    alpha, mu, sigma = inputs
    gaussians = tfd.MultivariateNormalDiag(loc = mu, scale_diag = sigma)
    alpha_prob = tfd.Categorical(probs = alpha)

    gm = tfd.MixtureSameFamily(
                mixture_distribution=alpha_prob,
                        components_distribution=gaussians)  
    Squeeze = Lambda(lambda x: K.squeeze(x, axis = 0))
    sample = gm.sample(1)
    return Squeeze(sample)

  def compute_output_shape(self, input_shape):
    return input_shape[1]

Наконец, я определяю всю модель:

def recurrent_density_network(seq_len, input_dim):
  x = Input(shape=(seq_len, input_dim,))

  alpha, mu, sigma = neural_net(x)
  mu, sigma = process_outputs(mu, sigma)

  gm = gaussian_mixture(alpha,mu,sigma)

  sample = GaussianMixture()([alpha, mu, sigma])    

  model = Model(inputs = x, outputs = sample)
  model.compile(optimizer = 'adam', loss = log_loss(gm))
  return model

со следующей нестандартной потерей:

def log_loss(gm, *args):
  def loss(y_true, sample):
    return -K.sum(gm.log_prob(y_true))
  return loss

Но я не понимаю проблемы, потому что мои входные и выходные данные имеют одинаковую форму.