是否有可能使用tensorflow的tpu_estimator()训练生成模型(即具有自定义损失计算的变分自动编码器)？

我的VAE的简化版本:

模型功能

def model_fn(features, labels, mode, params): 
  #Encoder layers  
  x = layers.Input()  
  h = conv1D()(x)  
  #BOTTLENECK LAYER  
  z_mean = Dense()(h)  
  z_log_var = Dense()(h)  
  def sampling(args):  
    z_mean_, z_log_var_ = args  
    epsilon = tf.random_normal()  
    return z_mean_ + tf.exp(z_log_var_/2)*epsilon  
  z = Lambda(sampling, name='lambda')([z_mean, z_log_var])
  #Decoder Layers
  h = Dense(z)
  x_decoded = TimeDistributed(Dense(activation='softmax'))(h)
  #VAE
  vae = tf.keras.models.Model(x, x_decoded)
  #VAE LOSS
  def vae_loss(x,x_decoded_mean):
    x = flatten(x)
    x_decoded_mean = flatten(x_decoded_mean)
    xent_loss = binary_crossentropy(x, x_decoded_mean)
    kl_loss = mean(1 + z_log_var - square(z_mean) - exp(z_log_var))
    return xent_loss …