from tools.collections import AttrDict

__C = AttrDict()

cfg = __C

__C.train_params = AttrDict()
__C.train_params.epochs = 300
__C.train_params.use_seed = False
__C.train_params.seed = 0

__C.optim = AttrDict()
__C.optim.init_lr = 0.025
__C.optim.min_lr = 1e-5
__C.optim.lr_schedule = 'cosine'  # cosine poly
__C.optim.momentum = 0.9
__C.optim.weight_decay = 4e-5
__C.optim.use_grad_clip = False
__C.optim.grad_clip = 10


__C.optim.label_smooth = False   # False, True
__C.optim.smooth_alpha = 0.25

__C.optim.auxiliary = True
__C.optim.auxiliary_weight = 0.4

__C.optim.if_resume = False
__C.optim.resume = AttrDict()
__C.optim.resume.load_path = ''
__C.optim.resume.load_epoch = 0

__C.data = AttrDict()
__C.data.num_workers = 16
__C.data.batch_size = 512
__C.data.dataset = 'imagenet' # imagenet
__C.data.train_data_type = 'lmdb'
__C.data.val_data_type = 'lmdb'
__C.data.patch_dataset = False
__C.data.num_examples = 100000
__C.data.input_size = (3, 64, 64)
__C.data.scaled_size = (3, 80, 80)
__C.data.type_of_data_aug = 'random_sized'  # random_sized / rand_scale
__C.data.random_sized = AttrDict()
__C.data.random_sized.min_scale = 0.08
__C.data.mean = [0.485, 0.456, 0.406]
__C.data.std = [0.229, 0.224, 0.225]
__C.data.color = True

__C.optim.use_warm_up = True    # False, True
__C.optim.warm_up = AttrDict()
__C.optim.warm_up.epoch = 5
__C.optim.warm_up.init_lr = 0.0001
__C.optim.warm_up.target_lr = 0.1

__C.optim.use_multi_stage = False
__C.optim.multi_stage = AttrDict()
__C.optim.multi_stage.stage_epochs = 330

__C.optim.cosine = AttrDict()
__C.optim.cosine.use_restart = False
__C.optim.cosine.restart=AttrDict()
__C.optim.cosine.restart.lr_period = [10, 20, 40, 80, 160, 320]
__C.optim.cosine.restart.lr_step = [0, 10, 30, 70, 150, 310]