b-schubert
commented
1 year ago Dear @bobosfw, this repo belongs to scCODA a statistical model for compositional data analysis of single-cell data. Though we are using TensorFlow-probability, we are not the package's maintainer. Please post your issue at their respective Github project.
When I try to train a very simple and small model with tensorflow and tensorflow_probability, the code is following
def neg_log_likelihood(y_true, y_hat): return -y_hat.log_prob(y_true)
inputs = tf.keras.layers.Input(shape=(X_train.shape[1], )) h1 = tf.keras.layers.BatchNormalization()(inputs) h2 = tf.keras.layers.Dense(4, activation=tf.nn.relu, name="layer_1")(h1) rate = tf.keras.layers.Dense(1+1, activation=tf.exp, name="layer_2")(h2) y = tfp.layers.DistributionLambda(lambda t: tfd.Skellam(rate1=t[..., 0:1], rate2=t[..., 1:]), name="prob")(rate)
model = tf.keras.Model(inputs=inputs, outputs=y) model.compile(tf.keras.optimizers.Adam(learning_rate=0.01), loss=neg_log_likelihood) history = model.fit(x=X_train, y=y_train, validation_data=(X_val, y_val), epochs=100, batch_size=128, shuffle=True)
When I start training the model, I get a tensorflow warning WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
Although I ignore the WARNING and keep training the model, I found the model training is very slow. The ETA is almost 2 min per epoch. The shape of the training set is (121562, 24)
So how can I fix the problem? thanks
tensorflow version--2.8.0 tensorflow_probability version--0.16.0