# Do we have a config file to parse?

args_config, remaining = config_parser.parse_known_args()

if args_config.config:

with open(args_config.config, 'r') as f:

cfg = yaml.safe_load(f)

parser.set_defaults(**cfg)

# The main arg parser parses the rest of the args, the usual

# defaults will have been overridden if config file specified.

args = parser.parse_args(remaining)

# Cache the args as a text string to save them in the output dir later

args_text = yaml.safe_dump(args.__dict__, default_flow_style=False)

setup_default_logging()

args, args_text = _parse_args()

dev_env = initialize_device(force_cpu=args.force_cpu, amp=args.amp, channels_last=args.channels_last)

if dev_env.distributed:

_logger.info('Training in distributed mode with multiple processes, 1 device per process. Process %d, total %d.'

% (dev_env.global_rank, dev_env.world_size))

else:

_logger.info('Training with a single process on 1 device.')

print("Update frequent = %d" % args.update_freq)

###############

assert not (args.scale_lr and args.scale_all_lr)

if args.scale_all_lr:

linear_scaled_lr = args.lr * args.batch_size * dev_env.world_size / 512.0

linear_scaled_warmup_lr = args.warmup_lr * args.batch_size * dev_env.world_size / 512.0

linear_scaled_min_lr = args.min_lr * args.batch_size * dev_env.world_size / 512.0

args.lr = linear_scaled_lr

args.warmup_lr = linear_scaled_warmup_lr

args.min_lr = linear_scaled_min_lr

elif args.scale_lr:

linear_scaled_lr = args.lr * args.batch_size * dev_env.world_size / 512.0

args.lr = linear_scaled_lr

###############

###########

random_seed(args.seed, 0) # Set all random seeds the same for model/state init (mandatory for XLA)

mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None

assert args.aug_splits == 0 or args.aug_splits > 1, 'A split of 1 makes no sense'

train_state = setup_train_task(args, dev_env, mixup_active)

train_cfg = train_state.train_cfg

# Set random seeds across ranks differently for train

# FIXME perhaps keep the same and just set diff seeds for dataloader worker process? what about TFDS?

random_seed(args.seed, dev_env.global_rank)

data_config, loader_eval, loader_train = setup_data(

args,

unwrap_model(train_state.model).default_cfg,

dev_env,

mixup_active)

#############

if args.step_in_iter:

train_state.lr_scheduler = convert_scheduler_from_epoch_to_iter(train_state.lr_scheduler,

len(loader_train) // args.update_freq)

############

# setup checkpoint manager

eval_metric = args.eval_metric

best_metric = None

best_epoch = None

checkpoint_manager = None

output_dir = None

if dev_env.primary:

if args.experiment:

exp_name = args.experiment

else:

exp_name = '-'.join([

datetime.now().strftime("%Y%m%d-%H%M%S"),

safe_model_name(args.model),

str(data_config['input_size'][-1])

])

output_dir = get_outdir(args.output if args.output else './output/train', exp_name)

checkpoint_manager = CheckpointManager(

hparams=vars(args),

checkpoint_dir=output_dir,

recovery_dir=output_dir,

metric_name=eval_metric,

metric_decreasing=True if eval_metric == 'loss' else False,

max_history=args.checkpoint_hist)

with open(os.path.join(output_dir, 'args.yaml'), 'w') as f:

f.write(args_text)

services = TrainServices(

monitor=Monitor(

output_dir=output_dir,

logger=_logger,

hparams=vars(args),

output_enabled=dev_env.primary),

checkpoint=checkpoint_manager,

)

try:

for epoch in range(train_state.epoch, train_cfg.num_epochs):

if dev_env.distributed and hasattr(loader_train.sampler, 'set_epoch'):

loader_train.sampler.set_epoch(epoch)

if args.mixup_off_epoch and epoch >= args.mixup_off_epoch:

if loader_train.mixup_enabled:

loader_train.mixup_enabled = False

train_metrics = train_one_epoch(

state=train_state,

services=services,

loader=loader_train,

dev_env=dev_env,

epoch=epoch,

update_freq=args.update_freq,

)

if dev_env.distributed and args.dist_bn in ('broadcast', 'reduce'):

if dev_env.primary:

_logger.info("Distributing BatchNorm running means and vars")

distribute_bn(train_state.model, args.dist_bn == 'reduce', dev_env)

eval_metrics = evaluate(

train_state.model,

train_state.eval_loss,

loader_eval,

services.monitor,

dev_env)

if train_state.model_ema is not None:

if dev_env.distributed and args.dist_bn in ('broadcast', 'reduce'):

distribute_bn(train_state.model_ema, args.dist_bn == 'reduce', dev_env)

ema_eval_metrics = evaluate(

train_state.model_ema.module,

train_state.eval_loss,

loader_eval,

services.monitor,

dev_env,

phase_suffix='EMA')

eval_metrics = ema_eval_metrics

if services.monitor is not None:

services.monitor.write_summary(

index=epoch,

results=dict(train=train_metrics, eval=eval_metrics))

if train_state.lr_scheduler is not None:

# step LR for next epoch

train_state.lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])

if checkpoint_manager is not None:

# save proper checkpoint with eval metric

best_checkpoint = checkpoint_manager.save_checkpoint(train_state,

save_epoch=(epoch + 1) % args.checkpoint_interval == 0,

metrics=eval_metrics)

best_metric, best_epoch = best_checkpoint.sort_key, best_checkpoint.epoch

train_state = replace(train_state, epoch=epoch + 1)

except KeyboardInterrupt:

pass

if best_metric is not None:

model = create_model(

args.model,

pretrained=args.pretrained,

num_classes=args.num_classes,

drop_rate=args.drop,

drop_connect_rate=args.drop_connect, # DEPRECATED, use drop_path

drop_path_rate=args.drop_path,

drop_block_rate=args.drop_block,

###########

layer_scale_init_value=args.layer_scale_init_value,

###########

global_pool=args.gp,

bn_tf=args.bn_tf,

bn_momentum=args.bn_momentum,

bn_eps=args.bn_eps,

scriptable=args.torchscript,

checkpoint_path=args.initial_checkpoint)

if args.num_classes is None:

assert hasattr(model, 'num_classes'), 'Model must have `num_classes` attr if not set on cmd line/config.'

args.num_classes = model.num_classes # FIXME handle model default vs config num_classes more elegantly

if dev_env.primary:

_logger.info(

f'Model {safe_model_name(args.model)} created, param count:{sum([m.numel() for m in model.parameters()])}')

# enable split bn (separate bn stats per batch-portion)

if args.split_bn:

assert args.aug_splits > 1

model = convert_splitbn_model(model, max(args.aug_splits, 2))

train_state = setup_model_and_optimizer(

dev_env=dev_env,

model=model,

optimizer=args.opt,

optimizer_cfg=optimizer_kwargs(cfg=args),

clip_fn=args.clip_mode if args.clip_grad is not None else None,

clip_value=args.clip_grad,

model_ema=args.model_ema,

model_ema_decay=args.model_ema_decay,

resume_path=args.resume,

resume_opt=not args.no_resume_opt,

use_syncbn=args.sync_bn,

)

# setup learning rate schedule and starting epoch

# FIXME move into updater?

lr_scheduler, num_epochs = create_scheduler(args, train_state.updater.optimizer)

if lr_scheduler is not None and train_state.epoch > 0:

lr_scheduler.step(train_state.epoch)

# setup loss function

if args.jsd_loss:

assert args.aug_splits > 1 # JSD only valid with aug splits set

train_loss_fn = JsdCrossEntropy(num_splits=args.aug_splits, smoothing=args.smoothing)

elif mixup_active:

# smoothing is handled with mixup target transform

if args.bce_loss:

train_loss_fn = BinaryCrossEntropy(target_threshold=args.bce_target_thresh)

else:

train_loss_fn = SoftTargetCrossEntropy()

elif args.smoothing:

if args.bce_loss:

train_loss_fn = BinaryCrossEntropy(smoothing=args.smoothing, target_threshold=args.bce_target_thresh)

else:

train_loss_fn = LabelSmoothingCrossEntropy(smoothing=args.smoothing)

else:

train_loss_fn = nn.CrossEntropyLoss()

eval_loss_fn = nn.CrossEntropyLoss()

dev_env.to_device(train_loss_fn, eval_loss_fn)

if dev_env.primary:

_logger.info('Scheduled epochs: {}'.format(num_epochs))

train_cfg = TrainCfg(

num_epochs=num_epochs,

log_interval=args.log_interval,

recovery_interval=args.recovery_interval,

)

train_state = replace(

train_state,

lr_scheduler=lr_scheduler,

train_loss=train_loss_fn,

eval_loss=eval_loss_fn,

train_cfg=train_cfg,

)

data_config = resolve_data_config(vars(args), default_cfg=default_cfg, verbose=dev_env.primary)

# create the train and eval datasets

dataset_train = create_dataset(

args.dataset,

root=args.data_dir, split=args.train_split, is_training=True,

batch_size=args.batch_size, repeats=args.epoch_repeats)

dataset_eval = create_dataset(

args.dataset,

root=args.data_dir, split=args.val_split, is_training=False, batch_size=args.batch_size)

_logger.info(f'len(dataset_train): {len(dataset_train)}')

_logger.info(f'len(datset_eval): {len(dataset_eval)}')

# setup mixup / cutmix

mixup_cfg = None

if mixup_active:

mixup_cfg = MixupCfg(

prob=args.mixup_prob, switch_prob=args.mixup_switch_prob, mode=args.mixup_mode,

mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, cutmix_minmax=args.cutmix_minmax,

label_smoothing=args.smoothing, num_classes=args.num_classes)

# wrap dataset in AugMix helper

if args.aug_splits > 1:

dataset_train = AugMixDataset(dataset_train, num_splits=args.aug_splits)

# create data loaders w/ augmentation pipeiine

train_interpolation = args.train_interpolation

if args.no_aug or not train_interpolation:

train_interpolation = data_config['interpolation']

if args.no_aug:

train_aug_cfg = None

else:

train_aug_cfg = AugCfg(

re_prob=args.reprob,

re_mode=args.remode,

re_count=args.recount,

ratio_range=args.ratio,

scale_range=args.scale,

hflip_prob=args.hflip,

vflip_prob=args.vflip,

color_jitter=args.color_jitter,

auto_augment=args.aa,

num_aug_splits=args.aug_splits,

)

train_pp_cfg = PreprocessCfg(

input_size=data_config['input_size'],

interpolation=train_interpolation,

crop_pct=data_config['crop_pct'],

mean=data_config['mean'],

std=data_config['std'],

aug=train_aug_cfg,

)

# if using PyTorch XLA and RandomErasing is enabled, we must normalize and do RE in transforms on CPU

normalize_in_transform = dev_env.type_xla and args.reprob > 0

dataset_train.transform = create_transform_v2(

cfg=train_pp_cfg, is_training=True, normalize=normalize_in_transform)

loader_train = create_loader_v2(

dataset_train,

batch_size=args.batch_size,

is_training=True,

normalize=not normalize_in_transform,

pp_cfg=train_pp_cfg,

mix_cfg=mixup_cfg,

num_aug_repeats=args.aug_repeats,

num_workers=args.workers,

pin_memory=args.pin_mem,

use_multi_epochs_loader=args.use_multi_epochs_loader

)

eval_pp_cfg = PreprocessCfg(

input_size=data_config['input_size'],

interpolation=data_config['interpolation'],

crop_pct=data_config['crop_pct'],

mean=data_config['mean'],

std=data_config['std'],

)

dataset_eval.transform = create_transform_v2(

cfg=eval_pp_cfg, is_training=False, normalize=normalize_in_transform)

eval_workers = args.workers

if 'tfds' in args.dataset:

# FIXME reduce validation issues when using TFDS w/ workers and distributed training

eval_workers = min(2, args.workers)

loader_eval = create_loader_v2(

dataset_eval,

batch_size=args.validation_batch_size or args.batch_size,

is_training=False,

normalize=not normalize_in_transform,

pp_cfg=eval_pp_cfg,

num_workers=eval_workers,

pin_memory=args.pin_mem,

)

state: TrainState,

services: TrainServices,

loader,

dev_env: DeviceEnv,

epoch: int,

update_freq: int,

):

tracker = Tracker()

loss_meter = AvgTensor() # FIXME move loss meter into task specific TaskMetric

state.model.train()

state.updater.reset() # zero-grad

step_end_idx = len(loader) - 1

tracker.mark_iter()

for step_idx, (sample, target) in enumerate(loader):

tracker.mark_iter_data_end()

# FIXME move forward + loss into model 'task' wrapper

with dev_env.autocast():

output = state.model(sample)

loss = state.train_loss(output, target)

#########

state.updater.apply(loss / update_freq, accumulate=False if (step_idx+1) % update_freq == 0 else True)

# state.updater.apply(loss)

#########

tracker.mark_iter_step_end()

state.updater.after_step(

after_train_step,

state,

services,

dev_env,

step_idx,

step_end_idx,

tracker,

loss_meter,

(output, target, loss),

update_freq,

)

tracker.mark_iter()

# end for

if hasattr(state.updater.optimizer, 'sync_lookahead'):

state.updater.optimizer.sync_lookahead()

state: TrainState,

services: TrainServices,

dev_env: DeviceEnv,

step_idx: int,

step_end_idx: int,

tracker: Tracker,

loss_meter: AvgTensor,

tensors: Tuple[torch.Tensor, ...],

update_freq: int,

):

"""

end_step = step_idx == step_end_idx

with torch.no_grad():

output, target, loss = tensors

loss_meter.update(loss, output.shape[0])

# if state.model_ema is not None:

if state.model_ema is not None and (step_idx+1) % update_freq == 0:

# FIXME should ema update be included here or in train / updater step? does it matter?

state.model_ema.update(state.model)

if (step_idx+1) % update_freq == 0:

# state = replace(state, step_count_global=state.step_count_global + 1)

state.step_count_global = state.step_count_global + 1

cfg = state.train_cfg

if services is not None and services.monitor is not None and end_step or (step_idx + 1) % cfg.log_interval == 0:

global_batch_size = dev_env.world_size * output.shape[0]

loss_avg = loss_meter.compute()

if services.monitor is not None:

lr_avg = state.updater.get_average_lr()

services.monitor.log_step(

'Train',

step=step_idx,

step_end=step_end_idx,

epoch=state.epoch,

loss=loss_avg.item(),

rate=tracker.get_avg_iter_rate(global_batch_size),

lr=lr_avg,

)

if services is not None and services.checkpoint is not None and cfg.recovery_interval and (

end_step or (step_idx + 1) % cfg.recovery_interval == 0):

services.checkpoint.save_recovery(state)

# if state.lr_scheduler is not None:

if state.lr_scheduler is not None and (step_idx + 1) % update_freq == 0:

# FIXME perform scheduler update here or via updater after_step call?

model: nn.Module,

loss_fn: nn.Module,

loader,

logger: Monitor,

dev_env: DeviceEnv,

phase_suffix: str = '',

log_interval: int = 10,

):

tracker = Tracker()

losses_m = AvgTensor()

accuracy_m = AccuracyTopK() # FIXME move loss and accuracy modules into task specific TaskMetric obj

model.eval()

end_idx = len(loader) - 1

tracker.mark_iter()

with torch.no_grad():

for step_idx, (sample, target) in enumerate(loader):

tracker.mark_iter_data_end()

last_step = step_idx == end_idx

with dev_env.autocast():

output = model(sample)

if isinstance(output, (tuple, list)):

output = output[0]

loss = loss_fn(output, target)

# FIXME, explictly marking step for XLA use since I'm not using the parallel xm loader

# need to investigate whether parallel loader wrapper is helpful on tpu-vm or only use for 2-vm setup.

if dev_env.type_xla:

dev_env.mark_step()

elif dev_env.type_cuda:

dev_env.synchronize()

# FIXME uncommenting this fixes race btw model `output`/`loss` and loss_m/accuracy_m meter input

# for PyTorch XLA GPU use.

# This issue does not exist for normal PyTorch w/ GPU (CUDA) or PyTorch XLA w/ TPU.

# loss.item()

tracker.mark_iter_step_end()

losses_m.update(loss, output.size(0))

accuracy_m.update(output, target)

if logger is not None and last_step or step_idx % log_interval == 0:

top1, top5 = accuracy_m.compute().values()

loss_avg = losses_m.compute()

logger.log_step(

'Eval',

step=step_idx,

step_end=end_idx,

loss=loss_avg.item(),

top1=top1.item(),

top5=top5.item(),

phase_suffix=phase_suffix,

)

tracker.mark_iter()

top1, top5 = accuracy_m.compute().values()

results = OrderedDict([(phase_suffix + 'loss', losses_m.compute().item()), (phase_suffix + 'top1', top1.item()), (phase_suffix + 'top5', top5.item())])