# Parameters for ModelTrainer:
# ==============================================================================
#
TrainConfig.input_size = 224
TrainConfig.use_swa = True
#
TrainConfig.drop_rate = 0.2
TrainConfig.feat_ext = True

# training parameters
TrainConfig.train_batch_size = 16
TrainConfig.test_batch_size = 16
TrainConfig.num_epochs = 100

# optimizer parameters
TrainConfig.lr_schedule_type = "reducelr_plateau"
TrainConfig.opt_type =  "lion"

#
TrainConfig.use_lookahead_opt = True
TrainConfig.weight_decay_rate = 0.01
TrainConfig.opt_beta1 = 0.9
TrainConfig.opt_beta2 = 0.999
TrainConfig.opt_eps = 1e-8  # 1e-6
TrainConfig.momentum = 0.9

# learning rate schedule parameters
TrainConfig.initial_lr = 1e-10
TrainConfig.max_lr = 1e-4   # 0.1
TrainConfig.final_lr = 1e-4
TrainConfig.final_lr_scale = 0.05

TrainConfig.lr_decay_rate = 0.3 #0.1
TrainConfig.warmup_proportion = 0.2
TrainConfig.num_warmup_steps = 0 #3125 #10000
TrainConfig.lr_scheduler_in_epoch = False
TrainConfig.step_size = 2


#
TrainConfig.use_clip_grad = True
TrainConfig.gradient_accumulation_steps = 1
TrainConfig.max_grad_norm = 1.0

##
TrainConfig.loss_func_type = 'cse'
TrainConfig.label_smooth_coef = 0.2
#
TrainConfig.random_seed = 12345
TrainConfig.use_amp = False
TrainConfig.show_per_step = 20
TrainConfig.total_steps = 0 #125000
TrainConfig.save_freq = 1000

# serialization
TrainConfig.out_dir = r'/home/uzver/Документы/code/naruto_sign/out/models/classification/naruto_sign/'