From 3966dd0cfd548c53fed6aae59601736739e30b9b Mon Sep 17 00:00:00 2001 From: SmilingWolf Date: Sat, 5 Mar 2022 16:23:28 +0100 Subject: [PATCH] train_convnext: switch to LAMB Can't get stable training with heavier augmentations with SGD no matter the learning rate --- train_convnext.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/train_convnext.py b/train_convnext.py index d3bdda3..10399fe 100644 --- a/train_convnext.py +++ b/train_convnext.py @@ -3,9 +3,9 @@ import numpy as np import tensorflow as tf import wandb -from tensorflow.keras.optimizers import SGD from tensorflow_addons.losses import SigmoidFocalCrossEntropy from tensorflow_addons.metrics import F1Score +from tensorflow_addons.optimizers import LAMB from wandb.keras import WandbCallback from Generator.ParseTFRecord import DataGenerator @@ -50,11 +50,11 @@ def scheduler(epoch, lr): global_batch_size = 32 * multiplier * strategy.num_replicas_in_sync # Training schedule - warmup_epochs = 5 + warmup_epochs = 0 total_epochs = 100 # Learning rate - max_learning_rate = 0.01 * multiplier + max_learning_rate = 0.0005 * multiplier warmup_learning_rate = max_learning_rate * 0.1 final_learning_rate = max_learning_rate * 0.01 @@ -131,7 +131,7 @@ def scheduler(epoch, lr): loss = SigmoidFocalCrossEntropy( reduction=tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE ) - opt = SGD(learning_rate=warmup_learning_rate, momentum=0.9, nesterov=True) + opt = LAMB(learning_rate=warmup_learning_rate) model.compile(optimizer=opt, loss=loss, metrics=[f1]) sched = tf.keras.callbacks.LearningRateScheduler(scheduler, verbose=True)