CODEWITHSUNDEEP
tensorflowkerasdeep-learning
DY92
DY92

Reputation: 437

Best practices in Tensorflow 2.0(Training step)

In tensorflow 2.0 you don't have to worry about training phase(batch size, number of epochs etc), because everything can be defined in compile method: model.fit(X_train,Y_train,batch_size = 64,epochs = 100).

But I have seen the following code style:

optimizer = tf.keras.optimizers.Adam(0.001)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()

@tf.function
def train_step(inputs, labels):
  with tf.GradientTape() as tape:
    predictions = model(inputs, training=True)
    regularization_loss = tf.math.add_n(model.losses)
    pred_loss = loss_fn(labels, predictions)
    total_loss = pred_loss + regularization_loss

  gradients = tape.gradient(total_loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))

for epoch in range(NUM_EPOCHS):
  for inputs, labels in train_data:
    train_step(inputs, labels)
  print("Finished epoch", epoch)

So here you can observe "more detailed" code, where you manually define by for loops you training procedure.

I have following question: what is the best practice in Tensorflow 2.0? I haven't found a any complete tutorial.

Upvotes: 0

Views: 397

Answers (1)

Daniel Möller
Daniel Möller

Reputation: 86650

Use what is best for your needs.

Both methods are documented in Tensorflow tutorials.

If you don't need anything special, no extra losses, strange metrics or intricate gradient computation, just use a model.fit() or a model.fit_generator(). This is totally ok and makes your life easier.

A custom training loop might come in handy when you have complicated models with non-trivial loss/gradients calculation.

Up to now, two applications I tried were easier with this:

  • Training a GAN's generator and discriminator simultaneously without having to do the generation step twice. (It's complicated because you have a loss function that applies to different y_true values, and each case should update only a part of the model) - The other option would require to have a few separate models, each model with its own trainable=True/False configuration, and train then in separate phases.
  • Training inputs (good for style transfer models) -- Alternatively, create a custom layer that takes dummy inputs and that outputs its own trainable weights. But it gets complicated to compile several loss functions for each of the outputs of the base and style networks.

Upvotes: 2

Related Questions