def data_creator(config):
# torch.manual_seed(args.seed + torch.distributed.get_rank())
args = config["args"]
train_dir = join(args.data, "train")
val_dir = join(args.data, "val")
if args.mock_data:
util.mock_data(train_dir, val_dir)
# todo: verbose should depend on rank
data_config = resolve_data_config(vars(args), verbose=True)
dataset_train = Dataset(join(args.data, "train"))
dataset_eval = Dataset(join(args.data, "val"))
collate_fn = None
if args.prefetcher and args.mixup > 0:
# collate conflict (need to support deinterleaving in collate mixup)
assert args.num_aug_splits == 0
collate_fn = FastCollateMixup(args.mixup, args.smoothing,
args.num_classes)
common_params = dict(
input_size=data_config["input_size"],
use_prefetcher=args.prefetcher,
mean=data_config["mean"],
std=data_config["std"],
num_workers=1,
distributed=args.distributed,
pin_memory=args.pin_mem)
train_loader = create_loader(
dataset_train,
is_training=True,
batch_size=config[BATCH_SIZE],
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
re_split=args.resplit,
collate_fn=collate_fn,
color_jitter=args.color_jitter,
auto_augment=args.aa,
interpolation=args.train_interpolation,
num_aug_splits=args.num_aug_splits, # always 0 right now
**common_params)
eval_loader = create_loader(
dataset_eval,
is_training=False,
batch_size=args.validation_batch_size_multiplier * config[BATCH_SIZE],
interpolation=data_config["interpolation"],
crop_pct=data_config["crop_pct"],
**common_params)
return train_loader, eval_loader
def main():
setup_default_logging()
args, args_text = _parse_args()
if args.log_wandb:
if has_wandb:
wandb.init(project=args.experiment, config=args)
else:
_logger.warning(
"You've requested to log metrics to wandb but package not found. "
"Metrics not being logged to wandb, try `pip install wandb`")
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
_logger.info(
'Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
_logger.info('Training with a single process on 1 GPUs.')
assert args.rank >= 0
# resolve AMP arguments based on PyTorch / Apex availability
use_amp = None
if args.amp:
# `--amp` chooses native amp before apex (APEX ver not actively maintained)
if has_native_amp:
args.native_amp = True
elif has_apex:
args.apex_amp = True
if args.apex_amp and has_apex:
use_amp = 'apex'
elif args.native_amp and has_native_amp:
use_amp = 'native'
elif args.apex_amp or args.native_amp:
_logger.warning(
"Neither APEX or native Torch AMP is available, using float32. "
"Install NVIDA apex or upgrade to PyTorch 1.6")
random_seed(args.seed, args.rank)
model_KD = None
if args.kd_model_path is not None:
model_KD = build_kd_model(args)
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,
global_pool=args.gp,
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 args.local_rank == 0:
_logger.info(
f'Model {safe_model_name(args.model)} created, param count:{sum([m.numel() for m in model.parameters()])}'
)
data_config = resolve_data_config(vars(args),
model=model,
verbose=args.local_rank == 0)
# setup augmentation batch splits for contrastive loss or split bn
num_aug_splits = 0
if args.aug_splits > 0:
assert args.aug_splits > 1, 'A split of 1 makes no sense'
num_aug_splits = args.aug_splits
# enable split bn (separate bn stats per batch-portion)
if args.split_bn:
assert num_aug_splits > 1 or args.resplit
model = convert_splitbn_model(model, max(num_aug_splits, 2))
# move model to GPU, enable channels last layout if set
model.cuda()
if args.channels_last:
model = model.to(memory_format=torch.channels_last)
# setup synchronized BatchNorm for distributed training
if args.distributed and args.sync_bn:
assert not args.split_bn
if has_apex and use_amp == 'apex':
# Apex SyncBN preferred unless native amp is activated
model = convert_syncbn_model(model)
else:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
if args.local_rank == 0:
_logger.info(
'Converted model to use Synchronized BatchNorm. WARNING: You may have issues if using '
'zero initialized BN layers (enabled by default for ResNets) while sync-bn enabled.'
)
if args.torchscript:
assert not use_amp == 'apex', 'Cannot use APEX AMP with torchscripted model'
assert not args.sync_bn, 'Cannot use SyncBatchNorm with torchscripted model'
model = torch.jit.script(model)
optimizer = create_optimizer_v2(model, **optimizer_kwargs(cfg=args))
# setup automatic mixed-precision (AMP) loss scaling and op casting
amp_autocast = suppress # do nothing
loss_scaler = None
if use_amp == 'apex':
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
loss_scaler = ApexScaler()
if args.local_rank == 0:
_logger.info('Using NVIDIA APEX AMP. Training in mixed precision.')
elif use_amp == 'native':
amp_autocast = torch.cuda.amp.autocast
loss_scaler = NativeScaler()
if args.local_rank == 0:
_logger.info(
'Using native Torch AMP. Training in mixed precision.')
else:
if args.local_rank == 0:
_logger.info('AMP not enabled. Training in float32.')
# optionally resume from a checkpoint
resume_epoch = None
if args.resume:
resume_epoch = resume_checkpoint(
model,
args.resume,
optimizer=None if args.no_resume_opt else optimizer,
loss_scaler=None if args.no_resume_opt else loss_scaler,
log_info=args.local_rank == 0)
# setup exponential moving average of model weights, SWA could be used here too
model_ema = None
if args.model_ema:
# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
model_ema = ModelEmaV2(
model,
decay=args.model_ema_decay,
device='cpu' if args.model_ema_force_cpu else None)
if args.resume:
load_checkpoint(model_ema.module, args.resume, use_ema=True)
# setup distributed training
if args.distributed:
if has_apex and use_amp == 'apex':
# Apex DDP preferred unless native amp is activated
if args.local_rank == 0:
_logger.info("Using NVIDIA APEX DistributedDataParallel.")
model = ApexDDP(model, delay_allreduce=True)
else:
if args.local_rank == 0:
_logger.info("Using native Torch DistributedDataParallel.")
model = NativeDDP(model,
device_ids=[args.local_rank],
broadcast_buffers=not args.no_ddp_bb)
# NOTE: EMA model does not need to be wrapped by DDP
# setup learning rate schedule and starting epoch
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
start_epoch = 0
if args.start_epoch is not None:
# a specified start_epoch will always override the resume epoch
start_epoch = args.start_epoch
elif resume_epoch is not None:
start_epoch = resume_epoch
if lr_scheduler is not None and start_epoch > 0:
lr_scheduler.step(start_epoch)
if args.local_rank == 0:
_logger.info('Scheduled epochs: {}'.format(num_epochs))
# create the train and eval datasets
dataset_train = create_dataset(args.dataset,
root=args.data_dir,
split=args.train_split,
is_training=True,
class_map=args.class_map,
download=args.dataset_download,
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,
class_map=args.class_map,
download=args.dataset_download,
batch_size=args.batch_size)
# setup mixup / cutmix
collate_fn = None
mixup_fn = None
mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
if mixup_active:
mixup_args = dict(mixup_alpha=args.mixup,
cutmix_alpha=args.cutmix,
cutmix_minmax=args.cutmix_minmax,
prob=args.mixup_prob,
switch_prob=args.mixup_switch_prob,
mode=args.mixup_mode,
label_smoothing=args.smoothing,
num_classes=args.num_classes)
if args.prefetcher:
assert not num_aug_splits # collate conflict (need to support deinterleaving in collate mixup)
collate_fn = FastCollateMixup(**mixup_args)
else:
mixup_fn = Mixup(**mixup_args)
# wrap dataset in AugMix helper
if num_aug_splits > 1:
dataset_train = AugMixDataset(dataset_train, num_splits=num_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']
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
no_aug=args.no_aug,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
re_split=args.resplit,
scale=args.scale,
ratio=args.ratio,
hflip=args.hflip,
vflip=args.vflip,
color_jitter=args.color_jitter,
auto_augment=args.aa,
num_aug_repeats=args.aug_repeats,
num_aug_splits=num_aug_splits,
interpolation=train_interpolation,
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
pin_memory=args.pin_mem,
use_multi_epochs_loader=args.use_multi_epochs_loader,
worker_seeding=args.worker_seeding,
)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=args.validation_batch_size or args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
crop_pct=data_config['crop_pct'],
pin_memory=args.pin_mem,
)
# setup loss function
if args.jsd_loss:
assert num_aug_splits > 1 # JSD only valid with aug splits set
train_loss_fn = JsdCrossEntropy(num_splits=num_aug_splits,
smoothing=args.smoothing)
elif mixup_active:
# smoothing is handled with mixup target transform which outputs sparse, soft targets
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()
train_loss_fn = train_loss_fn.cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
# setup checkpoint saver and eval metric tracking
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
output_dir = None
if args.rank == 0:
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)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(model=model,
optimizer=optimizer,
args=args,
model_ema=model_ema,
amp_scaler=loss_scaler,
checkpoint_dir=output_dir,
recovery_dir=output_dir,
decreasing=decreasing,
max_history=args.checkpoint_hist)
with open(os.path.join(output_dir, 'args.yaml'), 'w') as f:
f.write(args_text)
try:
for epoch in range(start_epoch, num_epochs):
if args.distributed and hasattr(loader_train.sampler, 'set_epoch'):
loader_train.sampler.set_epoch(epoch)
train_metrics = train_one_epoch(epoch,
model,
loader_train,
optimizer,
train_loss_fn,
args,
lr_scheduler=lr_scheduler,
saver=saver,
output_dir=output_dir,
amp_autocast=amp_autocast,
loss_scaler=loss_scaler,
model_ema=model_ema,
mixup_fn=mixup_fn,
model_KD=model_KD)
if args.distributed and args.dist_bn in ('broadcast', 'reduce'):
if args.local_rank == 0:
_logger.info(
"Distributing BatchNorm running means and vars")
distribute_bn(model, args.world_size, args.dist_bn == 'reduce')
eval_metrics = validate(model,
loader_eval,
validate_loss_fn,
args,
amp_autocast=amp_autocast)
if model_ema is not None and not args.model_ema_force_cpu:
if args.distributed and args.dist_bn in ('broadcast',
'reduce'):
distribute_bn(model_ema, args.world_size,
args.dist_bn == 'reduce')
ema_eval_metrics = validate(model_ema.module,
loader_eval,
validate_loss_fn,
args,
amp_autocast=amp_autocast,
log_suffix=' (EMA)')
eval_metrics = ema_eval_metrics
if lr_scheduler is not None:
# step LR for next epoch
lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
if output_dir is not None:
update_summary(epoch,
train_metrics,
eval_metrics,
os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None,
log_wandb=args.log_wandb and has_wandb)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
best_metric, best_epoch = saver.save_checkpoint(
epoch, metric=save_metric)
except KeyboardInterrupt:
pass
if best_metric is not None:
_logger.info('*** Best metric: {0} (epoch {1})'.format(
best_metric, best_epoch))
def main():
setup_default_logging()
args, args_text = _parse_args()
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
_logger.info('Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
_logger.info('Training with a single process on 1 GPUs.')
assert args.rank >= 0
# resolve AMP arguments based on PyTorch / Apex availability
use_amp = None
if args.amp:
# for backwards compat, `--amp` arg tries apex before native amp
if has_apex:
args.apex_amp = True
elif has_native_amp:
args.native_amp = True
if args.apex_amp and has_apex:
use_amp = 'apex'
elif args.native_amp and has_native_amp:
use_amp = 'native'
elif args.apex_amp or args.native_amp:
_logger.warning("Neither APEX or native Torch AMP is available, using float32. "
"Install NVIDA apex or upgrade to PyTorch 1.6")
torch.manual_seed(args.seed + args.rank)
####################################################################################
# Start - SparseML optional load weights from SparseZoo
####################################################################################
if args.initial_checkpoint == "zoo":
# Load checkpoint from base weights associated with given SparseZoo recipe
if args.sparseml_recipe.startswith("zoo:"):
args.initial_checkpoint = Zoo.download_recipe_base_framework_files(
args.sparseml_recipe,
extensions=[".pth.tar", ".pth"]
)[0]
else:
raise ValueError(
"Attempting to load weights from SparseZoo recipe, but not given a "
"SparseZoo recipe stub. When initial-checkpoint is set to 'zoo'. "
"sparseml-recipe must start with 'zoo:' and be a SparseZoo model "
f"stub. sparseml-recipe was set to {args.sparseml_recipe}"
)
elif args.initial_checkpoint.startswith("zoo:"):
# Load weights from a SparseZoo model stub
zoo_model = Zoo.load_model_from_stub(args.initial_checkpoint)
args.initial_checkpoint = zoo_model.download_framework_files(extensions=[".pth"])
####################################################################################
# End - SparseML optional load weights from SparseZoo
####################################################################################
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,
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 args.local_rank == 0:
_logger.info('Model %s created, param count: %d' %
(args.model, sum([m.numel() for m in model.parameters()])))
data_config = resolve_data_config(vars(args), model=model, verbose=args.local_rank == 0)
# setup augmentation batch splits for contrastive loss or split bn
num_aug_splits = 0
if args.aug_splits > 0:
assert args.aug_splits > 1, 'A split of 1 makes no sense'
num_aug_splits = args.aug_splits
# enable split bn (separate bn stats per batch-portion)
if args.split_bn:
assert num_aug_splits > 1 or args.resplit
model = convert_splitbn_model(model, max(num_aug_splits, 2))
# move model to GPU, enable channels last layout if set
model.cuda()
if args.channels_last:
model = model.to(memory_format=torch.channels_last)
# setup synchronized BatchNorm for distributed training
if args.distributed and args.sync_bn:
assert not args.split_bn
if has_apex and use_amp != 'native':
# Apex SyncBN preferred unless native amp is activated
model = convert_syncbn_model(model)
else:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
if args.local_rank == 0:
_logger.info(
'Converted model to use Synchronized BatchNorm. WARNING: You may have issues if using '
'zero initialized BN layers (enabled by default for ResNets) while sync-bn enabled.')
if args.torchscript:
assert not use_amp == 'apex', 'Cannot use APEX AMP with torchscripted model'
assert not args.sync_bn, 'Cannot use SyncBatchNorm with torchscripted model'
model = torch.jit.script(model)
optimizer = create_optimizer(args, model)
# setup automatic mixed-precision (AMP) loss scaling and op casting
amp_autocast = suppress # do nothing
loss_scaler = None
if use_amp == 'apex':
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
loss_scaler = ApexScaler()
if args.local_rank == 0:
_logger.info('Using NVIDIA APEX AMP. Training in mixed precision.')
elif use_amp == 'native':
amp_autocast = torch.cuda.amp.autocast
loss_scaler = NativeScaler()
if args.local_rank == 0:
_logger.info('Using native Torch AMP. Training in mixed precision.')
else:
if args.local_rank == 0:
_logger.info('AMP not enabled. Training in float32.')
# optionally resume from a checkpoint
resume_epoch = None
if args.resume:
resume_epoch = resume_checkpoint(
model, args.resume,
optimizer=None if args.no_resume_opt else optimizer,
loss_scaler=None if args.no_resume_opt else loss_scaler,
log_info=args.local_rank == 0)
# setup exponential moving average of model weights, SWA could be used here too
model_ema = None
if args.model_ema:
# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
model_ema = ModelEmaV2(
model, decay=args.model_ema_decay, device='cpu' if args.model_ema_force_cpu else None)
if args.resume:
load_checkpoint(model_ema.module, args.resume, use_ema=True)
# setup distributed training
if args.distributed:
if has_apex and use_amp != 'native':
# Apex DDP preferred unless native amp is activated
if args.local_rank == 0:
_logger.info("Using NVIDIA APEX DistributedDataParallel.")
model = ApexDDP(model, delay_allreduce=True)
else:
if args.local_rank == 0:
_logger.info("Using native Torch DistributedDataParallel.")
model = NativeDDP(model, device_ids=[args.local_rank]) # can use device str in Torch >= 1.1
# NOTE: EMA model does not need to be wrapped by DDP
# setup learning rate schedule and starting epoch
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
start_epoch = 0
if args.start_epoch is not None:
# a specified start_epoch will always override the resume epoch
start_epoch = args.start_epoch
elif resume_epoch is not None:
start_epoch = resume_epoch
if lr_scheduler is not None and start_epoch > 0:
lr_scheduler.step(start_epoch)
# 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)
dataset_eval = create_dataset(
args.dataset, root=args.data_dir, split=args.val_split, is_training=False, batch_size=args.batch_size)
# setup mixup / cutmix
collate_fn = None
mixup_fn = None
mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
if mixup_active:
mixup_args = dict(
mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, cutmix_minmax=args.cutmix_minmax,
prob=args.mixup_prob, switch_prob=args.mixup_switch_prob, mode=args.mixup_mode,
label_smoothing=args.smoothing, num_classes=args.num_classes)
if args.prefetcher:
assert not num_aug_splits # collate conflict (need to support deinterleaving in collate mixup)
collate_fn = FastCollateMixup(**mixup_args)
else:
mixup_fn = Mixup(**mixup_args)
# wrap dataset in AugMix helper
if num_aug_splits > 1:
dataset_train = AugMixDataset(dataset_train, num_splits=num_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']
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
no_aug=args.no_aug,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
re_split=args.resplit,
scale=args.scale,
ratio=args.ratio,
hflip=args.hflip,
vflip=args.vflip,
color_jitter=args.color_jitter,
auto_augment=args.aa,
num_aug_splits=num_aug_splits,
interpolation=train_interpolation,
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
pin_memory=args.pin_mem,
use_multi_epochs_loader=args.use_multi_epochs_loader
)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=args.validation_batch_size_multiplier * args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
crop_pct=data_config['crop_pct'],
pin_memory=args.pin_mem,
)
# setup loss function
if args.jsd:
assert num_aug_splits > 1 # JSD only valid with aug splits set
train_loss_fn = JsdCrossEntropy(num_splits=num_aug_splits, smoothing=args.smoothing).cuda()
elif mixup_active:
# smoothing is handled with mixup target transform
train_loss_fn = SoftTargetCrossEntropy().cuda()
elif args.smoothing:
train_loss_fn = LabelSmoothingCrossEntropy(smoothing=args.smoothing).cuda()
else:
train_loss_fn = nn.CrossEntropyLoss().cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
# setup checkpoint saver and eval metric tracking
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
output_dir = ''
if args.local_rank == 0:
output_base = args.output if args.output else './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"),
args.model,
str(data_config['input_size'][-1])
])
output_dir = get_outdir(output_base, 'train', exp_name)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(
model=model, optimizer=optimizer, args=args, model_ema=model_ema, amp_scaler=loss_scaler,
checkpoint_dir=output_dir, recovery_dir=output_dir, decreasing=decreasing, max_history=args.checkpoint_hist)
with open(os.path.join(output_dir, 'args.yaml'), 'w') as f:
f.write(args_text)
####################################################################################
# Start SparseML Integration
####################################################################################
sparseml_loggers = (
[PythonLogger(), TensorBoardLogger(log_path=output_dir)]
if output_dir
else None
)
manager = ScheduledModifierManager.from_yaml(args.sparseml_recipe)
optimizer = ScheduledOptimizer(
optimizer,
model,
manager,
steps_per_epoch=len(loader_train),
loggers=sparseml_loggers
)
# override lr scheduler if recipe makes any LR updates
if any("LearningRate" in str(modifier) for modifier in manager.modifiers):
_logger.info("Disabling timm LR scheduler, managing LR using SparseML recipe")
lr_scheduler = None
if manager.max_epochs:
_logger.info(
f"Overriding max_epochs to {manager.max_epochs} from SparseML recipe"
)
num_epochs = manager.max_epochs or num_epochs
####################################################################################
# End SparseML Integration
####################################################################################
if args.local_rank == 0:
_logger.info('Scheduled epochs: {}'.format(num_epochs))
try:
for epoch in range(start_epoch, num_epochs):
if args.distributed and hasattr(loader_train.sampler, 'set_epoch'):
loader_train.sampler.set_epoch(epoch)
train_metrics = train_one_epoch(
epoch, model, loader_train, optimizer, train_loss_fn, args,
lr_scheduler=lr_scheduler, saver=saver, output_dir=output_dir,
amp_autocast=amp_autocast, loss_scaler=loss_scaler, model_ema=model_ema, mixup_fn=mixup_fn)
if args.distributed and args.dist_bn in ('broadcast', 'reduce'):
if args.local_rank == 0:
_logger.info("Distributing BatchNorm running means and vars")
distribute_bn(model, args.world_size, args.dist_bn == 'reduce')
eval_metrics = validate(model, loader_eval, validate_loss_fn, args, amp_autocast=amp_autocast)
if model_ema is not None and not args.model_ema_force_cpu:
if args.distributed and args.dist_bn in ('broadcast', 'reduce'):
distribute_bn(model_ema, args.world_size, args.dist_bn == 'reduce')
ema_eval_metrics = validate(
model_ema.module, loader_eval, validate_loss_fn, args, amp_autocast=amp_autocast, log_suffix=' (EMA)')
eval_metrics = ema_eval_metrics
if lr_scheduler is not None:
# step LR for next epoch
lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
update_summary(
epoch, train_metrics, eval_metrics, os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
best_metric, best_epoch = saver.save_checkpoint(epoch, metric=save_metric)
#################################################################################
# Start SparseML ONNX Export
#################################################################################
if output_dir:
_logger.info(
f"training complete, exporting ONNX to {output_dir}/model.onnx"
)
exporter = ModuleExporter(model, output_dir)
exporter.export_onnx(torch.randn((1, *data_config["input_size"])))
#################################################################################
# End SparseML ONNX Export
#################################################################################
except KeyboardInterrupt:
pass
if best_metric is not None:
_logger.info('*** Best metric: {0} (epoch {1})'.format(best_metric, best_epoch))
def main():
args = parser.parse_args()
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
if args.distributed and args.num_gpu > 1:
print(
'Using more than one GPU per process in distributed mode is not allowed. Setting num_gpu to 1.'
)
args.num_gpu = 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.num_gpu = 1
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
assert args.rank >= 0
if args.distributed:
print(
'Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
print('Training with a single process on %d GPUs.' % args.num_gpu)
torch.manual_seed(args.seed + args.rank)
model = create_model(args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
drop_rate=args.drop,
global_pool=args.gp,
bn_tf=args.bn_tf,
bn_momentum=args.bn_momentum,
bn_eps=args.bn_eps,
checkpoint_path=args.initial_checkpoint)
print('Model %s created, param count: %d' %
(args.model, sum([m.numel() for m in model.parameters()])))
data_config = resolve_data_config(model,
args,
verbose=args.local_rank == 0)
# optionally resume from a checkpoint
start_epoch = 0
optimizer_state = None
if args.resume:
optimizer_state, start_epoch = resume_checkpoint(
model, args.resume, args.start_epoch)
if args.num_gpu > 1:
if args.amp:
print(
'Warning: AMP does not work well with nn.DataParallel, disabling. '
'Use distributed mode for multi-GPU AMP.')
args.amp = False
model = nn.DataParallel(model,
device_ids=list(range(args.num_gpu))).cuda()
else:
if args.distributed and args.sync_bn and has_apex:
model = convert_syncbn_model(model)
model.cuda()
optimizer = create_optimizer(args, model)
if optimizer_state is not None:
optimizer.load_state_dict(optimizer_state)
if has_apex and args.amp:
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
use_amp = True
print('AMP enabled')
else:
use_amp = False
print('AMP disabled')
model_ema = None
if args.model_ema:
model_ema = ModelEma(model,
decay=args.model_ema_decay,
device='cpu' if args.model_ema_force_cpu else '',
resume=args.resume)
if args.distributed:
model = DDP(model, delay_allreduce=True)
if model_ema is not None and not args.model_ema_force_cpu:
# must also distribute EMA model to allow validation
model_ema.ema = DDP(model_ema.ema, delay_allreduce=True)
model_ema.ema_has_module = True
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
if start_epoch > 0:
lr_scheduler.step(start_epoch)
if args.local_rank == 0:
print('Scheduled epochs: ', num_epochs)
train_dir = os.path.join(args.data, 'train')
if not os.path.exists(train_dir):
print('Error: training folder does not exist at: %s' % train_dir)
exit(1)
dataset_train = Dataset(train_dir)
collate_fn = None
if args.prefetcher and args.mixup > 0:
collate_fn = FastCollateMixup(args.mixup, args.smoothing,
args.num_classes)
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
rand_erase_prob=args.reprob,
rand_erase_mode=args.remode,
interpolation=
'random', # FIXME cleanly resolve this? data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
)
eval_dir = os.path.join(args.data, 'validation')
if not os.path.isdir(eval_dir):
print('Error: validation folder does not exist at: %s' % eval_dir)
exit(1)
dataset_eval = Dataset(eval_dir)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=4 * args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
)
if args.mixup > 0.:
# smoothing is handled with mixup label transform
train_loss_fn = SoftTargetCrossEntropy().cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
elif args.smoothing:
train_loss_fn = LabelSmoothingCrossEntropy(
smoothing=args.smoothing).cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
else:
train_loss_fn = nn.CrossEntropyLoss().cuda()
validate_loss_fn = train_loss_fn
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
output_dir = ''
if args.local_rank == 0:
output_base = args.output if args.output else './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"), args.model,
str(data_config['input_size'][-1])
])
output_dir = get_outdir(output_base, 'train', exp_name)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(checkpoint_dir=output_dir,
decreasing=decreasing)
try:
for epoch in range(start_epoch, num_epochs):
if args.distributed:
loader_train.sampler.set_epoch(epoch)
train_metrics = train_epoch(epoch,
model,
loader_train,
optimizer,
train_loss_fn,
args,
lr_scheduler=lr_scheduler,
saver=saver,
output_dir=output_dir,
use_amp=use_amp,
model_ema=model_ema)
eval_metrics = validate(model, loader_eval, validate_loss_fn, args)
if model_ema is not None and not args.model_ema_force_cpu:
ema_eval_metrics = validate(model_ema.ema,
loader_eval,
validate_loss_fn,
args,
log_suffix=' (EMA)')
eval_metrics = ema_eval_metrics
if lr_scheduler is not None:
lr_scheduler.step(epoch, eval_metrics[eval_metric])
update_summary(epoch,
train_metrics,
eval_metrics,
os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
best_metric, best_epoch = saver.save_checkpoint(
model,
optimizer,
args,
epoch=epoch + 1,
model_ema=model_ema,
metric=save_metric)
except KeyboardInterrupt:
pass
if best_metric is not None:
print('*** Best metric: {0} (epoch {1})'.format(
best_metric, best_epoch))
def main():
setup_default_logging()
args, args_text = _parse_args()
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
_logger.info(
'Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
_logger.info('Training with a single process on 1 GPUs.')
assert args.rank >= 0
# resolve AMP arguments based on PyTorch / Apex availability
use_amp = None
if args.amp:
# for backwards compat, `--amp` arg tries apex before native amp
if has_apex:
args.apex_amp = True
elif has_native_amp:
args.native_amp = True
if args.apex_amp and has_apex:
use_amp = 'apex'
elif args.native_amp and has_native_amp:
use_amp = 'native'
elif args.apex_amp or args.native_amp:
_logger.warning(
"Neither APEX or native Torch AMP is available, using float32. "
"Install NVIDA apex or upgrade to PyTorch 1.6")
torch.manual_seed(args.seed + args.rank)
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,
global_pool=args.gp,
bn_tf=args.bn_tf,
bn_momentum=args.bn_momentum,
bn_eps=args.bn_eps,
scriptable=args.torchscript,
use_cos_reg=args.cos_reg_component > 0,
checkpoint_path=args.initial_checkpoint)
with torch.cuda.device(0):
input = torch.randn(1, 3, 224, 224)
size_for_madd = 224 if args.img_size is None else args.img_size
# flops, params = get_model_complexity_info(model, (3, size_for_madd, size_for_madd), as_strings=True, print_per_layer_stat=True)
# print("=>Flops: " + flops)
# print("=>Params: " + params)
if args.local_rank == 0:
_logger.info('Model %s created, param count: %d' %
(args.model, sum([m.numel()
for m in model.parameters()])))
data_config = resolve_data_config(vars(args),
model=model,
verbose=args.local_rank == 0)
# setup augmentation batch splits for contrastive loss or split bn
num_aug_splits = 0
if args.aug_splits > 0:
assert args.aug_splits > 1, 'A split of 1 makes no sense'
num_aug_splits = args.aug_splits
# enable split bn (separate bn stats per batch-portion)
if args.split_bn:
assert num_aug_splits > 1 or args.resplit
model = convert_splitbn_model(model, max(num_aug_splits, 2))
# move model to GPU, enable channels last layout if set
model.cuda()
if args.channels_last:
model = model.to(memory_format=torch.channels_last)
# setup synchronized BatchNorm for distributed training
if args.distributed and args.sync_bn:
assert not args.split_bn
if has_apex and use_amp != 'native':
# Apex SyncBN preferred unless native amp is activated
model = convert_syncbn_model(model)
else:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
if args.local_rank == 0:
_logger.info(
'Converted model to use Synchronized BatchNorm. WARNING: You may have issues if using '
'zero initialized BN layers (enabled by default for ResNets) while sync-bn enabled.'
)
if args.torchscript:
assert not use_amp == 'apex', 'Cannot use APEX AMP with torchscripted model'
assert not args.sync_bn, 'Cannot use SyncBatchNorm with torchscripted model'
model = torch.jit.script(model)
optimizer = create_optimizer(args, model)
# setup automatic mixed-precision (AMP) loss scaling and op casting
amp_autocast = suppress # do nothing
loss_scaler = None
if use_amp == 'apex':
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
loss_scaler = ApexScaler()
if args.local_rank == 0:
_logger.info('Using NVIDIA APEX AMP. Training in mixed precision.')
elif use_amp == 'native':
amp_autocast = torch.cuda.amp.autocast
loss_scaler = NativeScaler()
if args.local_rank == 0:
_logger.info(
'Using native Torch AMP. Training in mixed precision.')
else:
if args.local_rank == 0:
_logger.info('AMP not enabled. Training in float32.')
# optionally resume from a checkpoint
resume_epoch = None
if args.resume:
resume_epoch = resume_checkpoint(
model,
args.resume,
optimizer=None if args.no_resume_opt else optimizer,
loss_scaler=None if args.no_resume_opt else loss_scaler,
log_info=args.local_rank == 0)
# setup exponential moving average of model weights, SWA could be used here too
model_ema = None
if args.model_ema:
# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
model_ema = ModelEmaV2(
model,
decay=args.model_ema_decay,
device='cpu' if args.model_ema_force_cpu else None)
if args.resume:
load_checkpoint(model_ema.module, args.resume, use_ema=True)
# setup distributed training
if args.distributed:
if has_apex and use_amp != 'native':
# Apex DDP preferred unless native amp is activated
if args.local_rank == 0:
_logger.info("Using NVIDIA APEX DistributedDataParallel.")
model = ApexDDP(model, delay_allreduce=True)
else:
if args.local_rank == 0:
_logger.info("Using native Torch DistributedDataParallel.")
model = NativeDDP(model, device_ids=[
args.local_rank
]) # can use device str in Torch >= 1.1
# NOTE: EMA model does not need to be wrapped by DDP
# setup learning rate schedule and starting epoch
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
start_epoch = 0
if args.start_epoch is not None:
# a specified start_epoch will always override the resume epoch
start_epoch = args.start_epoch
elif resume_epoch is not None:
start_epoch = resume_epoch
if lr_scheduler is not None and start_epoch > 0:
lr_scheduler.step(start_epoch)
if args.local_rank == 0:
_logger.info('Scheduled epochs: {}'.format(num_epochs))
# create the train and eval datasets
train_dir = os.path.join(args.data, 'train')
if not os.path.exists(train_dir):
_logger.error(
'Training folder does not exist at: {}'.format(train_dir))
exit(1)
if args.use_lmdb:
dataset_train = ImageFolderLMDB('../dataset_lmdb/train')
else:
dataset_train = Dataset(train_dir)
# dataset_train = Dataset(train_dir)
eval_dir = os.path.join(args.data, 'val')
if not os.path.isdir(eval_dir):
eval_dir = os.path.join(args.data, 'validation')
if not os.path.isdir(eval_dir):
_logger.error(
'Validation folder does not exist at: {}'.format(eval_dir))
exit(1)
if args.use_lmdb:
dataset_eval = ImageFolderLMDB('../dataset_lmdb/val')
else:
dataset_eval = Dataset(eval_dir)
# dataset_eval = Dataset(eval_dir)
# setup mixup / cutmix
collate_fn = None
mixup_fn = None
mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
if mixup_active:
mixup_args = dict(mixup_alpha=args.mixup,
cutmix_alpha=args.cutmix,
cutmix_minmax=args.cutmix_minmax,
prob=args.mixup_prob,
switch_prob=args.mixup_switch_prob,
mode=args.mixup_mode,
label_smoothing=args.smoothing,
num_classes=args.num_classes)
if args.prefetcher:
assert not num_aug_splits # collate conflict (need to support deinterleaving in collate mixup)
collate_fn = FastCollateMixup(**mixup_args)
else:
mixup_fn = Mixup(**mixup_args)
# wrap dataset in AugMix helper
if num_aug_splits > 1:
dataset_train = AugMixDataset(dataset_train, num_splits=num_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']
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
no_aug=args.no_aug,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
re_split=args.resplit,
scale=args.scale,
ratio=args.ratio,
hflip=args.hflip,
vflip=args.vflip,
color_jitter=args.color_jitter,
auto_augment=args.aa,
num_aug_splits=num_aug_splits,
interpolation=train_interpolation,
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
pin_memory=args.pin_mem,
use_multi_epochs_loader=args.use_multi_epochs_loader,
repeated_aug=args.use_repeated_aug,
world_size=args.world_size,
rank=args.rank)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=args.validation_batch_size_multiplier * args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
crop_pct=data_config['crop_pct'],
pin_memory=args.pin_mem,
)
loader_cali = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.cali_batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
no_aug=True,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
re_split=args.resplit,
scale=args.scale,
ratio=args.ratio,
hflip=args.hflip,
vflip=args.vflip,
color_jitter=args.color_jitter,
auto_augment=args.aa,
num_aug_splits=num_aug_splits,
interpolation=train_interpolation,
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=None,
pin_memory=args.pin_mem,
use_multi_epochs_loader=args.use_multi_epochs_loader,
repeated_aug=args.use_repeated_aug,
world_size=args.world_size,
rank=args.rank)
# setup loss function
if args.jsd:
assert num_aug_splits > 1 # JSD only valid with aug splits set
train_loss_fn = JsdCrossEntropy(num_splits=num_aug_splits,
smoothing=args.smoothing).cuda()
elif mixup_active:
# smoothing is handled with mixup target transform
if args.cos_reg_component > 0:
args.use_cos_reg_component = True
train_loss_fn = SoftTargetCrossEntropyCosReg(
n_comn=args.cos_reg_component).cuda()
else:
train_loss_fn = SoftTargetCrossEntropy().cuda()
args.use_cos_reg_component = False
elif args.smoothing:
train_loss_fn = LabelSmoothingCrossEntropy(
smoothing=args.smoothing).cuda()
else:
train_loss_fn = nn.CrossEntropyLoss().cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
# setup checkpoint saver and eval metric tracking
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
output_dir = ''
if args.local_rank == 0:
output_base = args.output if args.output else './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"), args.model,
str(data_config['input_size'][-1])
])
output_dir = get_outdir(output_base, 'train', exp_name)
code_dir = get_outdir(output_dir, 'code')
copy_tree(os.getcwd(), code_dir)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(model=model,
optimizer=optimizer,
args=args,
model_ema=model_ema,
amp_scaler=loss_scaler,
checkpoint_dir=output_dir,
recovery_dir=output_dir,
decreasing=decreasing)
with open(os.path.join(output_dir, 'args.yaml'), 'w') as f:
f.write(args_text)
try:
for epoch in range(start_epoch, num_epochs):
if args.distributed:
loader_train.sampler.set_epoch(epoch)
if not args.eval_only:
train_metrics = train_epoch(epoch,
model,
loader_train,
optimizer,
train_loss_fn,
args,
lr_scheduler=lr_scheduler,
saver=saver,
output_dir=output_dir,
amp_autocast=amp_autocast,
loss_scaler=loss_scaler,
model_ema=model_ema,
mixup_fn=mixup_fn)
if args.distributed and args.dist_bn in ('broadcast', 'reduce'):
if args.local_rank == 0:
_logger.info(
"Distributing BatchNorm running means and vars")
distribute_bn(model, args.world_size, args.dist_bn == 'reduce')
if args.max_iter > 0:
_ = validate(model,
loader_cali,
validate_loss_fn,
args,
amp_autocast=amp_autocast,
use_bn_calibration=True)
eval_metrics = validate(model,
loader_eval,
validate_loss_fn,
args,
amp_autocast=amp_autocast)
if model_ema is not None and not args.model_ema_force_cpu:
if args.distributed and args.dist_bn in ('broadcast',
'reduce'):
distribute_bn(model_ema, args.world_size,
args.dist_bn == 'reduce')
ema_eval_metrics = validate(model_ema.module,
loader_eval,
validate_loss_fn,
args,
amp_autocast=amp_autocast,
log_suffix=' (EMA)')
eval_metrics = ema_eval_metrics
if lr_scheduler is not None:
# step LR for next epoch
lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
if not args.eval_only:
update_summary(epoch,
train_metrics,
eval_metrics,
os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
best_metric, best_epoch = saver.save_checkpoint(
epoch, metric=save_metric)
if args.eval_only:
break
except KeyboardInterrupt:
pass
if best_metric is not None:
_logger.info('*** Best metric: {0} (epoch {1})'.format(
best_metric, best_epoch))
def main():
setup_default_logging()
args = parser.parse_args()
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
if args.distributed and args.num_gpu > 1:
logging.warning(
'Using more than one GPU per process in distributed mode is not allowed. Setting num_gpu to 1.'
)
args.num_gpu = 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.num_gpu = 1
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
assert args.rank >= 0
if args.distributed:
logging.info(
'Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
logging.info('Training with a single process on %d GPUs.' %
args.num_gpu)
torch.manual_seed(args.seed + args.rank)
model = create_model(args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
drop_rate=args.drop,
global_pool=args.gp,
bn_tf=args.bn_tf,
bn_momentum=args.bn_momentum,
bn_eps=args.bn_eps,
drop_connect_rate=0.2,
checkpoint_path=args.initial_checkpoint,
args=args)
flops, params = get_model_complexity_info(
model, (3, 224, 224),
as_strings=True,
print_per_layer_stat=args.display_info)
print('Flops: ' + flops)
print('Params: ' + params)
if args.KD_train:
teacher_model = create_model("efficientnet_b7_dq",
pretrained=True,
num_classes=args.num_classes,
drop_rate=args.drop,
global_pool=args.gp,
bn_tf=args.bn_tf,
bn_momentum=args.bn_momentum,
bn_eps=args.bn_eps,
drop_connect_rate=0.2,
checkpoint_path=args.initial_checkpoint,
args=args)
flops_teacher, params_teacher = get_model_complexity_info(
teacher_model, (3, 224, 224),
as_strings=True,
print_per_layer_stat=False)
print("Using KD training...")
print("FLOPs of teacher model: ", flops_teacher)
print("Params of teacher model: ", params_teacher)
if args.local_rank == 0:
logging.info('Model %s created, param count: %d' %
(args.model, sum([m.numel()
for m in model.parameters()])))
data_config = resolve_data_config(model,
args,
verbose=args.local_rank == 0)
# optionally resume from a checkpoint
start_epoch = 0
optimizer_state = None
if args.resume:
optimizer_state, start_epoch = resume_checkpoint(
model, args.resume, args.start_epoch)
# import pdb;pdb.set_trace()
if args.num_gpu > 1:
if args.amp:
logging.warning(
'AMP does not work well with nn.DataParallel, disabling. Use distributed mode for multi-GPU AMP.'
)
args.amp = False
model = nn.DataParallel(model,
device_ids=list(range(args.num_gpu))).cuda()
if args.KD_train:
teacher_model = nn.DataParallel(teacher_model,
device_ids=list(range(
args.num_gpu))).cuda()
else:
model.cuda()
if args.KD_train:
teacher_model.cuda()
optimizer = create_optimizer(args, model)
if optimizer_state is not None:
optimizer.load_state_dict(optimizer_state)
use_amp = False
if has_apex and args.amp:
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
use_amp = True
if args.local_rank == 0:
logging.info('NVIDIA APEX {}. AMP {}.'.format(
'installed' if has_apex else 'not installed',
'on' if use_amp else 'off'))
model_ema = None
if args.model_ema:
# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
# import pdb; pdb.set_trace()
model_ema = ModelEma(model,
decay=args.model_ema_decay,
device='cpu' if args.model_ema_force_cpu else '',
resume=args.resume)
if args.distributed:
if args.sync_bn:
try:
if has_apex:
model = convert_syncbn_model(model)
else:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
model)
if args.local_rank == 0:
logging.info(
'Converted model to use Synchronized BatchNorm.')
except Exception as e:
logging.error(
'Failed to enable Synchronized BatchNorm. Install Apex or Torch >= 1.1'
)
if has_apex:
model = DDP(model, delay_allreduce=True)
else:
if args.local_rank == 0:
logging.info(
"Using torch DistributedDataParallel. Install NVIDIA Apex for Apex DDP."
)
model = DDP(model,
device_ids=[args.local_rank
]) # can use device str in Torch >= 1.1
# NOTE: EMA model does not need to be wrapped by DDP
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
if start_epoch > 0:
lr_scheduler.step(start_epoch)
if args.local_rank == 0:
logging.info('Scheduled epochs: {}'.format(num_epochs))
train_dir = os.path.join(args.data, 'train')
if not os.path.exists(train_dir):
logging.error(
'Training folder does not exist at: {}'.format(train_dir))
exit(1)
dataset_train = Dataset(train_dir)
collate_fn = None
if args.prefetcher and args.mixup > 0:
collate_fn = FastCollateMixup(args.mixup, args.smoothing,
args.num_classes)
if args.auto_augment:
print('using auto data augumentation...')
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
rand_erase_prob=args.reprob,
rand_erase_mode=args.remode,
interpolation=
'bicubic', # FIXME cleanly resolve this? data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
use_auto_aug=args.auto_augment,
use_mixcut=args.mixcut,
)
eval_dir = os.path.join(args.data, 'val')
if not os.path.isdir(eval_dir):
logging.error(
'Validation folder does not exist at: {}'.format(eval_dir))
exit(1)
dataset_eval = Dataset(eval_dir)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=4 * args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
)
if args.mixup > 0.:
# smoothing is handled with mixup label transform
train_loss_fn = SoftTargetCrossEntropy().cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
elif args.smoothing:
train_loss_fn = LabelSmoothingCrossEntropy(
smoothing=args.smoothing).cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
else:
train_loss_fn = nn.CrossEntropyLoss().cuda()
validate_loss_fn = train_loss_fn
if args.KD_train:
train_loss_fn = nn.KLDivLoss(reduction='batchmean').cuda()
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
output_dir = ''
if args.local_rank == 0:
output_base = args.output if args.output else './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"), args.model,
str(data_config['input_size'][-1])
])
output_dir = get_outdir(output_base, 'train', exp_name)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(checkpoint_dir=output_dir,
decreasing=decreasing)
try:
# import pdb;pdb.set_trace()
for epoch in range(start_epoch, num_epochs):
if args.distributed:
loader_train.sampler.set_epoch(epoch)
# import pdb; pdb.set_trace()
if args.KD_train:
train_metrics = train_epoch(epoch,
model,
loader_train,
optimizer,
train_loss_fn,
args,
lr_scheduler=lr_scheduler,
saver=saver,
output_dir=output_dir,
use_amp=use_amp,
model_ema=model_ema,
teacher_model=teacher_model)
else:
train_metrics = train_epoch(epoch,
model,
loader_train,
optimizer,
train_loss_fn,
args,
lr_scheduler=lr_scheduler,
saver=saver,
output_dir=output_dir,
use_amp=use_amp,
model_ema=model_ema)
# def __init__(self, model, bits_activations=8, bits_parameters=8, bits_accum=32,
# overrides=None, mode=LinearQuantMode.SYMMETRIC, clip_acts=ClipMode.NONE,
# per_channel_wts=False, model_activation_stats=None, fp16=False, clip_n_stds=None,
# scale_approx_mult_bits=None):
# import distiller
# import pdb; pdb.set_trace()
# quantizer = quantization.PostTrainLinearQuantizer.from_args(model, args)
# quantizer.prepare_model(distiller.get_dummy_input(input_shape=model.input_shape))
# quantizer = distiller.quantization.PostTrainLinearQuantizer(model, bits_activations=8, bits_parameters=8)
# quantizer.prepare_model()
# distiller.utils.assign_layer_fq_names(model)
# # msglogger.info("Generating quantization calibration stats based on {0} users".format(args.qe_calibration))
# collector = distiller.data_loggers.QuantCalibrationStatsCollector(model)
# with collector_context(collector):
# eval_metrics = validate(model, loader_eval, validate_loss_fn, args)
# # Here call your model evaluation function, making sure to execute only
# # the portion of the dataset specified by the qe_calibration argument
# yaml_path = './dir/quantization_stats.yaml'
# collector.save(yaml_path)
eval_metrics = validate(model, loader_eval, validate_loss_fn, args)
if model_ema is not None and not args.model_ema_force_cpu:
ema_eval_metrics = validate(model_ema.ema,
loader_eval,
validate_loss_fn,
args,
log_suffix=' (EMA)')
eval_metrics = ema_eval_metrics
if lr_scheduler is not None:
lr_scheduler.step(epoch, eval_metrics[eval_metric])
update_summary(epoch,
train_metrics,
eval_metrics,
os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
best_metric, best_epoch = saver.save_checkpoint(
model,
optimizer,
args,
epoch=epoch + 1,
model_ema=model_ema,
metric=save_metric)
except KeyboardInterrupt:
pass
if best_metric is not None:
logging.info('*** Best metric: {0} (epoch {1})'.format(
best_metric, best_epoch))
def main():
get_logger("./")
args = parser.parse_args()
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
if args.distributed and args.num_gpu > 1:
logging.warning(
'Using more than one GPU per process in distributed mode is not allowed. Setting num_gpu to 1.'
)
args.num_gpu = 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.num_gpu = 1
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
assert args.rank >= 0
if args.distributed:
logging.info(
'Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
logging.info('Training with a single process on %d GPUs.' %
args.num_gpu)
logging.info("Exponential : {}".format(args.model_ema_decay))
logging.info("Color Jitter : {}".format(args.color_jitter))
logging.info("Model EMA Decay : {}".format(args.model_ema_decay))
torch.manual_seed(args.seed + args.rank)
model = eval(args.model)(config_path=args.config_path,
target_flops=args.target_flops,
num_classes=args.num_classes,
bn_momentum=args.bn_momentum,
activation=args.activation,
se=args.se)
if os.path.exists(args.initial_checkpoint):
load_checkpoint(model, args.initial_checkpoint)
if args.local_rank == 0:
logging.info('Model %s created, param count: %d' %
(args.model, sum([m.numel()
for m in model.parameters()])))
data_config = resolve_data_config(vars(args),
model=model,
verbose=args.local_rank == 0)
# optionally resume from a checkpoint
optimizer_state = None
resume_epoch = None
if args.resume:
optimizer_state, resume_epoch = resume_checkpoint(model, args.resume)
if args.num_gpu > 1:
if args.amp:
logging.warning(
'AMP does not work well with nn.DataParallel, disabling. Use distributed mode for multi-GPU AMP.'
)
args.amp = False
model = nn.DataParallel(model,
device_ids=list(range(args.num_gpu))).cuda()
else:
model.cuda()
logging.info(args.weight_decay)
optimizer = create_optimizer(args, model)
if optimizer_state is not None:
optimizer.load_state_dict(optimizer_state["optimizer"])
use_amp = False
if has_apex and args.amp:
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
use_amp = True
if args.local_rank == 0:
logging.info('NVIDIA APEX {}. AMP {}.'.format(
'installed' if has_apex else 'not installed',
'on' if use_amp else 'off'))
model_ema = None
if args.model_ema:
# Important to create EMA model after cuda(), DP wrapper, and AMP but
# before SyncBN and DDP wrapper
model_ema = ModelEma(model,
decay=args.model_ema_decay,
device='cpu' if args.model_ema_force_cpu else '',
resume=args.resume)
if args.distributed:
if args.sync_bn:
try:
if has_apex:
model = convert_syncbn_model(model)
else:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
model)
if args.local_rank == 0:
logging.info(
'Converted model to use Synchronized BatchNorm.')
except Exception as e:
logging.error(
'Failed to enable Synchronized BatchNorm. Install Apex or Torch >= 1.1'
)
if has_apex:
model = DDP(model, delay_allreduce=True)
else:
if args.local_rank == 0:
logging.info(
"Using torch DistributedDataParallel. Install NVIDIA Apex for Apex DDP."
)
# can use device str in Torch >= 1.1
model = DDP(model, device_ids=[args.local_rank])
# NOTE: EMA model does not need to be wrapped by DDP
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
start_epoch = 0
if args.start_epoch is not None:
# a specified start_epoch will always override the resume epoch
start_epoch = args.start_epoch
elif resume_epoch is not None:
start_epoch = resume_epoch
if start_epoch > 0:
lr_scheduler.step(start_epoch)
if args.local_rank == 0:
logging.info('Scheduled epochs: {}'.format(num_epochs))
if args.lmdb:
train_dir = os.path.join(args.data, 'train_lmdb', 'train.lmdb')
dataset_train = ImageFolderLMDB(train_dir, None, None)
else:
train_dir = os.path.join(args.data, 'train')
dataset_train = Dataset(train_dir)
collate_fn = None
if args.prefetcher and args.mixup > 0:
collate_fn = FastCollateMixup(args.mixup, args.smoothing,
args.num_classes)
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
rand_erase_prob=args.reprob,
rand_erase_mode=args.remode,
color_jitter=args.color_jitter,
interpolation='random',
# FIXME cleanly resolve this? data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
)
if args.lmdb:
eval_dir = os.path.join(args.data, 'test_lmdb', 'test.lmdb')
dataset_eval = ImageFolderLMDB(eval_dir, None, None)
else:
eval_dir = os.path.join(args.data, 'val')
dataset_eval = Dataset(eval_dir)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=4 * args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
)
if args.mixup > 0.:
# smoothing is handled with mixup label transform
train_loss_fn = SoftTargetCrossEntropy().cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
elif args.smoothing:
train_loss_fn = LabelSmoothingCrossEntropy(
smoothing=args.smoothing).cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
else:
train_loss_fn = nn.CrossEntropyLoss().cuda()
validate_loss_fn = train_loss_fn
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
output_dir = ''
if args.local_rank == 0:
output_base = args.output if args.output else './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"), args.model,
str(data_config['input_size'][-1])
])
output_dir = get_outdir(output_base, 'train', exp_name)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(checkpoint_dir=output_dir,
decreasing=decreasing)
try:
for epoch in range(start_epoch, num_epochs):
if args.distributed:
loader_train.sampler.set_epoch(epoch)
train_metrics = train_epoch(epoch,
model,
loader_train,
optimizer,
train_loss_fn,
args,
lr_scheduler=lr_scheduler,
saver=saver,
output_dir=output_dir,
use_amp=use_amp,
model_ema=model_ema)
eval_metrics = validate(model, loader_eval, validate_loss_fn, args)
if model_ema is not None and not args.model_ema_force_cpu:
ema_eval_metrics = validate(model_ema.ema,
loader_eval,
validate_loss_fn,
args,
log_suffix=' (EMA)')
eval_metrics = ema_eval_metrics
if lr_scheduler is not None:
# step LR for next epoch
lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
update_summary(epoch,
train_metrics,
eval_metrics,
os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
best_metric, best_epoch = saver.save_checkpoint(
model,
optimizer,
args,
epoch=epoch,
model_ema=model_ema,
metric=save_metric)
except KeyboardInterrupt:
pass
if best_metric is not None:
logging.info('*** Best metric: {0} (epoch {1})'.format(
best_metric, best_epoch))
def main():
setup_default_logging()
args, args_text = _parse_args()
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
if args.distributed and args.num_gpu > 1:
logging.warning(
'Using more than one GPU per process in distributed mode is not allowed. Setting num_gpu to 1.'
)
args.num_gpu = 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.num_gpu = 1
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
assert args.rank >= 0
DistributedManager.set_args(args)
if args.distributed:
logging.info(
'Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
logging.info('Training with a single process on %d GPUs.' %
args.num_gpu)
torch.manual_seed(args.seed + args.rank)
model = create_model(args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
drop_rate=args.drop,
drop_connect_rate=args.drop_connect,
drop_path_rate=args.drop_path,
drop_block_rate=args.drop_block,
global_pool=args.gp,
bn_tf=args.bn_tf,
bn_momentum=args.bn_momentum,
bn_eps=args.bn_eps,
checkpoint_path=args.initial_checkpoint)
if args.initial_checkpoint_pruned:
try:
data_config = resolve_data_config(vars(args),
model=model,
verbose=args.local_rank == 0)
model2 = load_module_from_ckpt(
model,
args.initial_checkpoint_pruned,
input_size=data_config['input_size'][1])
logging.info("New pruned model adapted from the checkpoint")
except Exception as e:
raise RuntimeError(e)
else:
model2 = model
if args.local_rank == 0:
logging.info('Model %s created, param count: %d' %
(args.model, sum([m.numel()
for m in model2.parameters()])))
data_config = resolve_data_config(vars(args),
model=model,
verbose=args.local_rank == 0)
num_aug_splits = 0
if args.aug_splits > 0:
assert args.aug_splits > 1, 'A split of 1 makes no sense'
num_aug_splits = args.aug_splits
if args.split_bn:
assert num_aug_splits > 1 or args.resplit
model = convert_splitbn_model(model, max(num_aug_splits, 2))
if args.num_gpu > 1:
model2 = nn.DataParallel(model2,
device_ids=list(range(args.num_gpu))).cuda()
else:
model2.cuda()
use_amp = False
if args.distributed:
model2 = nn.parallel.distributed.DistributedDataParallel(
model2,
device_ids=[args.local_rank]) # can use device str in Torch >= 1.1
# NOTE: EMA model does not need to be wrapped by DDP
train_dir = os.path.join(args.data, 'train')
if not os.path.exists(train_dir):
logging.error(
'Training folder does not exist at: {}'.format(train_dir))
exit(1)
dataset_train = Dataset(train_dir)
collate_fn = None
if args.prefetcher and args.mixup > 0:
collate_fn = FastCollateMixup(args.mixup, args.smoothing,
args.num_classes)
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
re_split=args.resplit,
color_jitter=args.color_jitter,
auto_augment=args.aa,
num_aug_splits=num_aug_splits,
interpolation=args.train_interpolation,
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
pin_memory=args.pin_mem,
)
eval_dir = os.path.join(args.data, 'val')
if not os.path.isdir(eval_dir):
eval_dir = os.path.join(args.data, 'validation')
if not os.path.isdir(eval_dir):
eval_dir = os.path.join(args.data, 'test')
if not os.path.isdir(eval_dir):
logging.error(
'Validation folder does not exist at: {}'.format(eval_dir))
exit(1)
test_dir = os.path.join(args.data, 'test')
if not os.path.isdir(test_dir):
test_dir = os.path.join(args.data, 'validation')
if not os.path.isdir(test_dir):
test_dir = os.path.join(args.data, 'val')
if not os.path.isdir(test_dir):
logging.error(
'Test folder does not exist at: {}'.format(test_dir))
exit(1)
dataset_eval = Dataset(eval_dir)
if args.prune_test:
dataset_test = Dataset(test_dir)
else:
dataset_test = Dataset(train_dir)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=args.validation_batch_size_multiplier * args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
crop_pct=data_config['crop_pct'],
pin_memory=args.pin_mem,
)
len_loader = int(
len(loader_eval) * (4 * args.batch_size) / args.batch_size_prune)
if args.prune_test:
len_loader = None
if args.prune:
loader_p = create_loader(
dataset_test,
input_size=data_config['input_size'],
batch_size=args.batch_size_prune,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
crop_pct=data_config['crop_pct'],
pin_memory=args.pin_mem,
)
if 'resnet' in model2.__class__.__name__.lower() or (
hasattr(model2, 'module')
and 'resnet' in model2.module.__class__.__name__.lower()):
list_channel_to_prune = compute_num_channels_per_layer_taylor(
model2,
data_config['input_size'],
loader_p,
pruning_ratio=args.pruning_ratio,
taylor_file=args.taylor_file,
local_rank=args.local_rank,
len_data_loader=len_loader,
prune_skip=args.prune_skip,
taylor_abs=args.taylor_abs,
prune_conv1=args.prune_conv1,
use_time=args.use_time,
distributed=args.distributed)
new_net = redesign_module_resnet(
model2,
list_channel_to_prune,
use_amp=use_amp,
distributed=args.distributed,
local_rank=args.local_rank,
input_size=data_config['input_size'][1])
else:
list_channel_to_prune = compute_num_channels_per_layer_taylor(
model2,
data_config['input_size'],
loader_p,
pruning_ratio=args.pruning_ratio,
taylor_file=args.taylor_file,
local_rank=args.local_rank,
len_data_loader=len_loader,
prune_pwl=not args.no_pwl,
taylor_abs=args.taylor_abs,
use_se=not args.use_eca,
use_time=args.use_time,
distributed=args.distributed)
new_net = redesign_module_efnet(
model2,
list_channel_to_prune,
use_amp=use_amp,
distributed=args.distributed,
local_rank=args.local_rank,
input_size=data_config['input_size'][1],
use_se=not args.use_eca)
new_net.train()
model.train()
if isinstance(model, nn.DataParallel) or isinstance(model, DDP):
model = model.module
else:
model = model.cuda()
co_mod = build_co_train_model(
model,
new_net.module.cpu() if hasattr(new_net, 'module') else new_net,
gamma=args.gamma_knowledge,
only_last=args.only_last,
progressive_IKD_factor=args.progressive_IKD_factor)
optimizer = create_optimizer(args, co_mod)
del model
del new_net
gc.collect()
torch.cuda.empty_cache()
if args.num_gpu > 1:
if args.amp:
logging.warning(
'AMP does not work well with nn.DataParallel, disabling. Use distributed mode for multi-GPU AMP.'
)
args.amp = False
co_mod = nn.DataParallel(co_mod,
device_ids=list(range(
args.num_gpu))).cuda()
else:
co_mod = co_mod.cuda()
use_amp = False
if has_apex and args.amp:
co_mod, optimizer = amp.initialize(co_mod,
optimizer,
opt_level='O1')
use_amp = True
if args.local_rank == 0:
logging.info('NVIDIA APEX {}. AMP {}.'.format(
'installed' if has_apex else 'not installed',
'on' if use_amp else 'off'))
if args.distributed:
if args.sync_bn:
try:
if has_apex and use_amp:
co_mod = convert_syncbn_model(co_mod)
else:
co_mod = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
co_mod)
if args.local_rank == 0:
logging.info(
'Converted model to use Synchronized BatchNorm.')
except Exception as e:
logging.error(
'Failed to enable Synchronized BatchNorm. Install Apex or Torch >= 1.1'
)
if has_apex and use_amp:
co_mod = DDP(co_mod, delay_allreduce=False)
else:
if args.local_rank == 0 and use_amp:
logging.info(
"Using torch DistributedDataParallel. Install NVIDIA Apex for Apex DDP."
)
co_mod = nn.parallel.distributed.DistributedDataParallel(
co_mod,
device_ids=[args.local_rank
]) # can use device str in Torch >= 1.1
# NOTE: EMA model does not need to be wrapped by DDP
co_mod.train()
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
start_epoch = 0
if args.start_epoch is not None:
# a specified start_epoch will always override the resume epoch
start_epoch = args.start_epoch
if lr_scheduler is not None and start_epoch > 0:
lr_scheduler.step(start_epoch)
if args.jsd:
assert num_aug_splits > 1 # JSD only valid with aug splits set
train_loss_fn = JsdCrossEntropy(num_splits=num_aug_splits,
smoothing=args.smoothing).cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
elif args.mixup > 0.:
# smoothing is handled with mixup label transform
train_loss_fn = SoftTargetCrossEntropy().cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
elif args.smoothing:
train_loss_fn = LabelSmoothingCrossEntropy(
smoothing=args.smoothing).cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
else:
train_loss_fn = nn.CrossEntropyLoss().cuda()
validate_loss_fn = train_loss_fn
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
output_dir = ''
if args.local_rank == 0:
output_base = args.output if args.output else './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"), args.model,
str(data_config['input_size'][-1])
])
output_dir = get_outdir(output_base, 'train', exp_name)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(checkpoint_dir=output_dir,
decreasing=decreasing)
with open(os.path.join(output_dir, 'args.yaml'), 'w') as f:
f.write(args_text)
try:
if args.local_rank == 0:
logging.info(f'First validation')
co_mod.eval()
eval_metrics = validate(co_mod, loader_eval, validate_loss_fn, args)
if args.local_rank == 0:
logging.info(f'Prec@top1 : {eval_metrics["prec1"]}')
co_mod.train()
for epoch in range(start_epoch, num_epochs):
torch.cuda.empty_cache()
if args.distributed:
loader_train.sampler.set_epoch(epoch)
train_metrics = train_epoch(epoch,
co_mod,
loader_train,
optimizer,
train_loss_fn,
args,
lr_scheduler=lr_scheduler,
saver=saver,
output_dir=output_dir,
use_amp=use_amp,
model_ema=None)
torch.cuda.empty_cache()
eval_metrics = validate(co_mod, loader_eval, validate_loss_fn,
args)
if lr_scheduler is not None:
# step LR for next epoch
lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
update_summary(epoch,
train_metrics,
eval_metrics,
os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
best_metric, best_epoch = saver.save_checkpoint(
co_mod,
optimizer,
args,
epoch=epoch,
model_ema=None,
metric=save_metric,
use_amp=use_amp)
except KeyboardInterrupt:
pass
if best_metric is not None:
logging.info('*** Best metric: {0} (epoch {1})'.format(
best_metric, best_epoch))
def train_imagenet_dq():
setup_default_logging()
args = parser.parse_args()
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
if args.distributed and args.num_gpu > 1:
logging.warning('Using more than one GPU per process in distributed mode is not allowed. Setting num_gpu to 1.')
args.num_gpu = 1
args.device = xm.xla_device()
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.num_gpu = 1
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
assert args.rank >= 0
if args.distributed:
logging.info('Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
logging.info('Training with a single process on %d GPUs.' % args.num_gpu)
torch.manual_seed(args.seed + args.rank)
device = xm.xla_device()
model = create_model(
args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
drop_rate=args.drop,
global_pool=args.gp,
bn_tf=args.bn_tf,
bn_momentum=args.bn_momentum,
bn_eps=args.bn_eps,
drop_connect_rate=0.2,
checkpoint_path=args.initial_checkpoint,
args = args).to(device)
flops, params = get_model_complexity_info(model, (3, 224, 224), as_strings=True, print_per_layer_stat=args.display_info)
print('Flops: ' + flops)
print('Params: ' + params)
if args.KD_train:
teacher_model = create_model(
"efficientnet_b7_dq",
pretrained=True,
num_classes=args.num_classes,
drop_rate=args.drop,
global_pool=args.gp,
bn_tf=args.bn_tf,
bn_momentum=args.bn_momentum,
bn_eps=args.bn_eps,
drop_connect_rate=0.2,
checkpoint_path=args.initial_checkpoint,
args = args)
flops_teacher, params_teacher = get_model_complexity_info(teacher_model, (3, 224, 224), as_strings=True, print_per_layer_stat=False)
print("Using KD training...")
print("FLOPs of teacher model: ", flops_teacher)
print("Params of teacher model: ", params_teacher)
if args.local_rank == 0:
logging.info('Model %s created, param count: %d' %
(args.model, sum([m.numel() for m in model.parameters()])))
data_config = resolve_data_config(model, args, verbose=args.local_rank == 0)
# optionally resume from a checkpoint
start_epoch = 0
optimizer_state = None
if args.resume:
optimizer_state, start_epoch = resume_checkpoint(model, args.resume, args.start_epoch)
# import pdb;pdb.set_trace()
torch.manual_seed(42)
if args.num_gpu > 1:
if args.amp:
logging.warning(
'AMP does not work well with nn.DataParallel, disabling. Use distributed mode for multi-GPU AMP.')
args.amp = False
# device = xm.xla_device()
# devices = (
# xm.get_xla_supported_devices(
# max_devices=num_cores) if num_cores != 0 else [])
# model = nn.DataParallel(model, device_ids=devices).cuda()
# model = model.to(device)
if args.KD_train:
teacher_model = nn.DataParallel(teacher_model, device_ids=list(range(args.num_gpu))).cuda()
else:
# device = xm.xla_device()
# model = model.to(device)
if args.KD_train:
teacher_model.cuda()
optimizer = create_optimizer(args, model)
if optimizer_state is not None:
optimizer.load_state_dict(optimizer_state)
use_amp = False
if has_apex and args.amp:
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
use_amp = True
if args.local_rank == 0:
logging.info('NVIDIA APEX {}. AMP {}.'.format(
'installed' if has_apex else 'not installed', 'on' if use_amp else 'off'))
model_ema = None
if args.model_ema:
# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
# import pdb; pdb.set_trace()
model_e = create_model(
args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
drop_rate=args.drop,
global_pool=args.gp,
bn_tf=args.bn_tf,
bn_momentum=args.bn_momentum,
bn_eps=args.bn_eps,
drop_connect_rate=0.2,
checkpoint_path=args.initial_checkpoint,
args = args).to(device)
model_ema = ModelEma(
model_e,
decay=args.model_ema_decay,
device='cpu' if args.model_ema_force_cpu else '',
resume=args.resume)
if args.distributed:
if args.sync_bn:
try:
if has_apex:
model = convert_syncbn_model(model)
else:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
if args.local_rank == 0:
logging.info('Converted model to use Synchronized BatchNorm.')
except Exception as e:
logging.error('Failed to enable Synchronized BatchNorm. Install Apex or Torch >= 1.1')
if has_apex:
model = DDP(model, delay_allreduce=True)
else:
if args.local_rank == 0:
logging.info("Using torch DistributedDataParallel. Install NVIDIA Apex for Apex DDP.")
model = DDP(model, device_ids=[args.local_rank]) # can use device str in Torch >= 1.1
# NOTE: EMA model does not need to be wrapped by DDP
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
if start_epoch > 0:
lr_scheduler.step(start_epoch)
if args.local_rank == 0:
logging.info('Scheduled epochs: {}'.format(num_epochs))
train_dir = os.path.join(args.data, 'train')
if not os.path.exists(train_dir):
logging.error('Training folder does not exist at: {}'.format(train_dir))
exit(1)
dataset_train = Dataset(train_dir)
collate_fn = None
if args.prefetcher and args.mixup > 0:
collate_fn = FastCollateMixup(args.mixup, args.smoothing, args.num_classes)
if args.auto_augment:
print('using auto data augumentation...')
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
rand_erase_prob=args.reprob,
rand_erase_mode=args.remode,
interpolation='bicubic', # FIXME cleanly resolve this? data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
use_auto_aug=args.auto_augment,
use_mixcut=args.mixcut,
)
eval_dir = os.path.join(args.data, 'val')
if not os.path.isdir(eval_dir):
logging.error('Validation folder does not exist at: {}'.format(eval_dir))
exit(1)
dataset_eval = Dataset(eval_dir)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size = args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
)
if args.mixup > 0.:
# smoothing is handled with mixup label transform
train_loss_fn = SoftTargetCrossEntropy()
validate_loss_fn = nn.CrossEntropyLoss()
elif args.smoothing:
train_loss_fn = LabelSmoothingCrossEntropy(smoothing=args.smoothing)
validate_loss_fn = nn.CrossEntropyLoss()
else:
train_loss_fn = nn.CrossEntropyLoss()
validate_loss_fn = train_loss_fn
if args.KD_train:
train_loss_fn = nn.KLDivLoss(reduction='batchmean')
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
output_dir = ''
if args.local_rank == 0:
output_base = args.output if args.output else './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"),
args.model,
str(data_config['input_size'][-1])
])
output_dir = get_outdir(output_base, 'train', exp_name)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(checkpoint_dir=output_dir, decreasing=decreasing)
def train_epoch(
epoch, model, loader, optimizer, loss_fn, args,
lr_scheduler=None, saver=None, output_dir='', use_amp=False, model_ema=None, teacher_model = None, loader_len=0):
if args.prefetcher and args.mixup > 0 and loader.mixup_enabled:
if args.mixup_off_epoch and epoch >= args.mixup_off_epoch:
loader.mixup_enabled = False
batch_time_m = AverageMeter()
data_time_m = AverageMeter()
losses_m = AverageMeter()
model.train()
if args.KD_train:
teacher_model.eval()
end = time.time()
last_idx = loader_len - 1
num_updates = epoch * loader_len
for batch_idx, (input, target) in loader:
last_batch = batch_idx == last_idx
data_time_m.update(time.time() - end)
if not args.prefetcher:
# input = input.cuda()
# target = target.cuda()
if args.mixup > 0.:
lam = 1.
if not args.mixup_off_epoch or epoch < args.mixup_off_epoch:
lam = np.random.beta(args.mixup, args.mixup)
input.mul_(lam).add_(1 - lam, input.flip(0))
target = mixup_target(target, args.num_classes, lam, args.smoothing)
r = np.random.rand(1)
if args.beta > 0 and r < args.cutmix_prob:
# generate mixed sample
lam = np.random.beta(args.beta, args.beta)
rand_index = torch.randperm(input.size()[0])
target_a = target
target_b = target[rand_index]
bbx1, bby1, bbx2, bby2 = rand_bbox(input.size(), lam)
input[:, :, bbx1:bbx2, bby1:bby2] = input[rand_index, :, bbx1:bbx2, bby1:bby2]
# adjust lambda to exactly match pixel ratio
lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (input.size()[-1] * input.size()[-2]))
# compute output
input_var = torch.autograd.Variable(input, requires_grad=True)
target_a_var = torch.autograd.Variable(target_a)
target_b_var = torch.autograd.Variable(target_b)
output = model(input_var)
loss = loss_fn(output, target_a_var) * lam + loss_fn(output, target_b_var) * (1. - lam)
else:
# NOTE KD Train is exclusive with mixcut, FIX it later
output = model(input)
if args.KD_train:
# teacher_model.cuda()
teacher_outputs_tmp = []
assert(input.shape[0]%args.teacher_step == 0)
step_size = int(input.shape[0]//args.teacher_step)
with torch.no_grad():
for k in range(0,int(input.shape[0]),step_size):
input_tmp = input[k:k+step_size,:,:,:]
teacher_outputs_tmp.append(teacher_model(input_tmp))
# torch.cuda.empty_cache()
# import pdb; pdb.set_trace()
teacher_outputs = torch.cat(teacher_outputs_tmp)
alpha = args.KD_alpha
T = args.KD_temperature
loss = loss_fn(F.log_softmax(output/T, dim=1),
F.softmax(teacher_outputs/T, dim=1)) * (alpha * T * T) + \
F.cross_entropy(output, target) * (1. - alpha)
else:
loss = loss_fn(output, target)
if not args.distributed:
losses_m.update(loss.item(), input.size(0))
optimizer.zero_grad()
if use_amp:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
#optimizer.step()
xm.optimizer_step(optimizer)
# torch.cuda.synchronize()
if model_ema is not None:
model_ema.update(model)
num_updates += 1
batch_time_m.update(time.time() - end)
if last_batch or batch_idx % args.log_interval == 0:
lrl = [param_group['lr'] for param_group in optimizer.param_groups]
lr = sum(lrl) / len(lrl)
if args.distributed:
reduced_loss = reduce_tensor(loss.data, args.world_size)
losses_m.update(reduced_loss.item(), input.size(0))
if args.local_rank == 0:
logging.info(
'Train: {} [{:>4d}/{} ({:>3.0f}%)] '
'Loss: {loss.val:>9.6f} ({loss.avg:>6.4f}) '
'Time: {batch_time.val:.3f}s, {rate:>7.2f}/s '
'({batch_time.avg:.3f}s, {rate_avg:>7.2f}/s) '
'LR: {lr:.3e} '
'Data: {data_time.val:.3f} ({data_time.avg:.3f})'.format(
epoch,
batch_idx, loader_len,
100. * batch_idx / last_idx,
loss=losses_m,
batch_time=batch_time_m,
rate=input.size(0) * args.world_size / batch_time_m.val,
rate_avg=input.size(0) * args.world_size / batch_time_m.avg,
lr=lr,
data_time=data_time_m))
if args.save_images and output_dir:
torchvision.utils.save_image(
input,
os.path.join(output_dir, 'train-batch-%d.jpg' % batch_idx),
padding=0,
normalize=True)
if saver is not None and args.recovery_interval and (
last_batch or (batch_idx + 1) % args.recovery_interval == 0):
save_epoch = epoch + 1 if last_batch else epoch
saver.save_recovery(
model, optimizer, args, save_epoch, model_ema=model_ema, batch_idx=batch_idx)
if lr_scheduler is not None:
lr_scheduler.step_update(num_updates=num_updates, metric=losses_m.avg)
end = time.time()
return OrderedDict([('loss', losses_m.avg)])
def validate(model, loader, loss_fn, args, log_suffix='',loader_len=0):
batch_time_m = AverageMeter()
losses_m = AverageMeter()
prec1_m = AverageMeter()
prec5_m = AverageMeter()
model.eval()
end = time.time()
last_idx = loader_len - 1
with torch.no_grad():
for batch_idx, (input, target) in loader:
last_batch = batch_idx == last_idx
# if not args.prefetcher:
# input = input.cuda()
# target = target.cuda()
output = model(input)
if isinstance(output, (tuple, list)):
output = output[0]
# augmentation reduction
reduce_factor = args.tta
if reduce_factor > 1:
output = output.unfold(0, reduce_factor, reduce_factor).mean(dim=2)
target = target[0:target.size(0):reduce_factor]
loss = loss_fn(output, target)
prec1, prec5 = accuracy(output, target, topk=(1, 5))
if args.distributed:
reduced_loss = reduce_tensor(loss.data, args.world_size)
prec1 = reduce_tensor(prec1, args.world_size)
prec5 = reduce_tensor(prec5, args.world_size)
else:
reduced_loss = loss.data
# torch.cuda.synchronize()
losses_m.update(reduced_loss.item(), input.size(0))
prec1_m.update(prec1.item(), output.size(0))
prec5_m.update(prec5.item(), output.size(0))
batch_time_m.update(time.time() - end)
end = time.time()
if args.local_rank == 0 and (last_batch or batch_idx % args.log_interval == 0):
log_name = 'Test' + log_suffix
logging.info(
'{0}: [{1:>4d}/{2}] '
'Time: {batch_time.val:.3f} ({batch_time.avg:.3f}) '
'Loss: {loss.val:>7.4f} ({loss.avg:>6.4f}) '
'Prec@1: {top1.val:>7.4f} ({top1.avg:>7.4f}) '
'Prec@5: {top5.val:>7.4f} ({top5.avg:>7.4f})'.format(
log_name, batch_idx, last_idx,
batch_time=batch_time_m, loss=losses_m,
top1=prec1_m, top5=prec5_m))
metrics = OrderedDict([('loss', losses_m.avg), ('prec1', prec1_m.avg), ('prec5', prec5_m.avg)])
return metrics
try:
# import pdb;pdb.set_trace()
for epoch in range(start_epoch, num_epochs):
loader_len=len(loader_train)
if args.distributed:
loader_train.sampler.set_epoch(epoch)
# import pdb; pdb.set_trace()
if args.KD_train:
train_metrics = train_epoch(
epoch, model, loader_train, optimizer, train_loss_fn, args,
lr_scheduler=lr_scheduler, saver=saver, output_dir=output_dir,
use_amp=use_amp, model_ema=model_ema, teacher_model = teacher_model)
else:
para_loader = dp.ParallelLoader(loader_train, [device])
train_metrics = train_epoch(
epoch, model, para_loader.per_device_loader(device), optimizer, train_loss_fn, args,
lr_scheduler=lr_scheduler, saver=saver, output_dir=output_dir,
use_amp=use_amp, model_ema=model_ema, loader_len=loader_len)
# def __init__(self, model, bits_activations=8, bits_parameters=8, bits_accum=32,
# overrides=None, mode=LinearQuantMode.SYMMETRIC, clip_acts=ClipMode.NONE,
# per_channel_wts=False, model_activation_stats=None, fp16=False, clip_n_stds=None,
# scale_approx_mult_bits=None):
# import distiller
# import pdb; pdb.set_trace()
# quantizer = quantization.PostTrainLinearQuantizer.from_args(model, args)
# quantizer.prepare_model(distiller.get_dummy_input(input_shape=model.input_shape))
# quantizer = distiller.quantization.PostTrainLinearQuantizer(model, bits_activations=8, bits_parameters=8)
# quantizer.prepare_model()
# distiller.utils.assign_layer_fq_names(model)
# # msglogger.info("Generating quantization calibration stats based on {0} users".format(args.qe_calibration))
# collector = distiller.data_loggers.QuantCalibrationStatsCollector(model)
# with collector_context(collector):
# eval_metrics = validate(model, loader_eval, validate_loss_fn, args)
# # Here call your model evaluation function, making sure to execute only
# # the portion of the dataset specified by the qe_calibration argument
# yaml_path = './dir/quantization_stats.yaml'
# collector.save(yaml_path)
loader_len_val = len(loader_eval)
para_loader = dp.ParallelLoader(loader_eval, [device])
eval_metrics = validate(model, para_loader.per_device_loader(device), validate_loss_fn, args, loader_len=loader_len_val)
if model_ema is not None and not args.model_ema_force_cpu:
ema_eval_metrics = validate(model_ema.ema, loader_eval, validate_loss_fn, args, log_suffix=' (EMA)')
eval_metrics = ema_eval_metrics
if lr_scheduler is not None:
lr_scheduler.step(epoch, eval_metrics[eval_metric])
update_summary(
epoch, train_metrics, eval_metrics, os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
best_metric, best_epoch = saver.save_checkpoint(
model, optimizer, args,
epoch=epoch + 1,
model_ema=model_ema,
metric=save_metric)
except KeyboardInterrupt:
pass
if best_metric is not None:
logging.info('*** Best metric: {0} (epoch {1})'.format(best_metric, best_epoch))
def main():
setup_default_logging()
args, args_text = _parse_args()
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
_logger.info('Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
_logger.info('Training with a single process on 1 GPUs.')
assert args.rank >= 0
if args.control_amp == 'amp':
args.amp = True
elif args.control_amp == 'apex':
args.apex_amp = True
elif args.control_amp == 'native':
args.native_amp = True
# resolve AMP arguments based on PyTorch / Apex availability
use_amp = None
if args.amp:
# for backwards compat, `--amp` arg tries apex before native amp
if has_apex:
args.apex_amp = True
elif has_native_amp:
args.native_amp = True
if args.apex_amp and has_apex:
use_amp = 'apex'
elif args.native_amp and has_native_amp:
use_amp = 'native'
elif args.apex_amp or args.native_amp:
_logger.warning("Neither APEX or native Torch AMP is available, using float32. "
"Install NVIDA apex or upgrade to PyTorch 1.6")
_logger.info(
'====================\n\n'
'Actfun: {}\n'
'LR: {}\n'
'Epochs: {}\n'
'p: {}\n'
'k: {}\n'
'g: {}\n'
'Extra channel multiplier: {}\n'
'AMP: {}\n'
'Weight Init: {}\n'
'\n===================='.format(args.actfun, args.lr, args.epochs, args.p, args.k, args.g,
args.extra_channel_mult, use_amp, args.weight_init))
torch.manual_seed(args.seed + args.rank)
model = create_model(
args.model,
pretrained=args.pretrained,
actfun=args.actfun,
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,
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,
p=args.p,
k=args.k,
g=args.g,
extra_channel_mult=args.extra_channel_mult,
weight_init_name=args.weight_init,
partial_ho_actfun=args.partial_ho_actfun
)
if args.tl:
if args.data == 'caltech101' and not os.path.exists('caltech101'):
dir_root = r'101_ObjectCategories'
dir_new = r'caltech101'
dir_new_train = os.path.join(dir_new, 'train')
dir_new_val = os.path.join(dir_new, 'val')
dir_new_test = os.path.join(dir_new, 'test')
if not os.path.exists(dir_new):
os.mkdir(dir_new)
os.mkdir(dir_new_train)
os.mkdir(dir_new_val)
os.mkdir(dir_new_test)
for dir2 in os.listdir(dir_root):
if dir2 != 'BACKGROUND_Google':
curr_path = os.path.join(dir_root, dir2)
new_path_train = os.path.join(dir_new_train, dir2)
new_path_val = os.path.join(dir_new_val, dir2)
new_path_test = os.path.join(dir_new_test, dir2)
if not os.path.exists(new_path_train):
os.mkdir(new_path_train)
if not os.path.exists(new_path_val):
os.mkdir(new_path_val)
if not os.path.exists(new_path_test):
os.mkdir(new_path_test)
train_upper = int(0.8 * len(os.listdir(curr_path)))
val_upper = int(0.9 * len(os.listdir(curr_path)))
curr_files_all = os.listdir(curr_path)
curr_files_train = curr_files_all[:train_upper]
curr_files_val = curr_files_all[train_upper:val_upper]
curr_files_test = curr_files_all[val_upper:]
for file in curr_files_train:
copyfile(os.path.join(curr_path, file),
os.path.join(new_path_train, file))
for file in curr_files_val:
copyfile(os.path.join(curr_path, file),
os.path.join(new_path_val, file))
for file in curr_files_test:
copyfile(os.path.join(curr_path, file),
os.path.join(new_path_test, file))
time.sleep(5)
if args.tl:
pre_model = create_model(
args.model,
pretrained=True,
actfun='swish',
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,
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,
p=args.p,
k=args.k,
g=args.g,
extra_channel_mult=args.extra_channel_mult,
weight_init_name=args.weight_init,
partial_ho_actfun=args.partial_ho_actfun
)
model = MLP.MLP(actfun=args.actfun,
input_dim=1280,
output_dim=args.num_classes,
k=args.k,
p=args.p,
g=args.g,
num_params=400_000,
permute_type='shuffle')
pre_model_layers = list(pre_model.children())
pre_model = torch.nn.Sequential(*pre_model_layers[:-1])
else:
pre_model = None
if args.local_rank == 0:
_logger.info('Model %s created, param count: %d' %
(args.model, sum([m.numel() for m in model.parameters()])))
data_config = resolve_data_config(vars(args), model=model, verbose=args.local_rank == 0)
# setup augmentation batch splits for contrastive loss or split bn
num_aug_splits = 0
if args.aug_splits > 0:
assert args.aug_splits > 1, 'A split of 1 makes no sense'
num_aug_splits = args.aug_splits
# enable split bn (separate bn stats per batch-portion)
if args.split_bn:
assert num_aug_splits > 1 or args.resplit
model = convert_splitbn_model(model, max(num_aug_splits, 2))
# move model to GPU, enable channels last layout if set
model.cuda()
if args.tl:
pre_model.cuda()
if args.channels_last:
model = model.to(memory_format=torch.channels_last)
# setup synchronized BatchNorm for distributed training
if args.distributed and args.sync_bn:
assert not args.split_bn
if has_apex and use_amp != 'native':
# Apex SyncBN preferred unless native amp is activated
model = convert_syncbn_model(model)
else:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
if args.local_rank == 0:
_logger.info(
'Converted model to use Synchronized BatchNorm. WARNING: You may have issues if using '
'zero initialized BN layers (enabled by default for ResNets) while sync-bn enabled.')
if args.torchscript:
assert not use_amp == 'apex', 'Cannot use APEX AMP with torchscripted model'
assert not args.sync_bn, 'Cannot use SyncBatchNorm with torchscripted model'
model = torch.jit.script(model)
if args.tl:
optimizer = torch.optim.Adam(model.parameters(), weight_decay=1e-5)
else:
optimizer = create_optimizer(args, model)
# setup automatic mixed-precision (AMP) loss scaling and op casting
amp_autocast = suppress # do nothing
loss_scaler = None
if use_amp == 'apex':
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
loss_scaler = ApexScaler()
if args.local_rank == 0:
_logger.info('Using NVIDIA APEX AMP. Training in mixed precision.')
elif use_amp == 'native':
amp_autocast = torch.cuda.amp.autocast
loss_scaler = NativeScaler()
if args.local_rank == 0:
_logger.info('Using native Torch AMP. Training in mixed precision.')
else:
if args.local_rank == 0:
_logger.info('AMP not enabled. Training in float32.')
if args.local_rank == 0:
_logger.info('\n--------------------\nModel:\n' + repr(model) + '--------------------')
# optionally resume from a checkpoint
resume_epoch = None
resume_path = os.path.join(args.resume, 'recover.pth.tar')
if args.resume and os.path.exists(resume_path):
resume_epoch = resume_checkpoint(
model, resume_path,
optimizer=None if args.no_resume_opt else optimizer,
loss_scaler=None if args.no_resume_opt else loss_scaler,
log_info=args.local_rank == 0)
cp_loaded = None
resume_epoch = None
checkname = 'recover'
if args.actfun != 'swish':
checkname = '{}_'.format(args.actfun) + checkname
check_path = os.path.join(args.check_path, checkname) + '.pth'
loader = None
if os.path.isfile(check_path):
loader = check_path
elif args.load_path != '' and os.path.isfile(args.load_path):
loader = args.load_path
if loader is not None:
cp_loaded = torch.load(loader)
model.load_state_dict(cp_loaded['model'])
optimizer.load_state_dict(cp_loaded['optimizer'])
resume_epoch = cp_loaded['epoch']
model.cuda()
loss_scaler.load_state_dict(cp_loaded['amp'])
if args.channels_last:
model = model.to(memory_format=torch.channels_last)
_logger.info('============ LOADED CHECKPOINT: Epoch {}'.format(resume_epoch))
model_raw = model
# setup exponential moving average of model weights, SWA could be used here too
model_ema = None
if args.model_ema:
# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
model_ema = ModelEmaV2(
model, decay=args.model_ema_decay, device='cpu' if args.model_ema_force_cpu else None)
if args.resume and os.path.exists(resume_path):
load_checkpoint(model_ema.module, args.resume, use_ema=True)
if cp_loaded is not None:
model_ema.load_state_dict(cp_loaded['model_ema'])
# setup distributed training
if args.distributed:
if has_apex and use_amp != 'native':
# Apex DDP preferred unless native amp is activated
if args.local_rank == 0:
_logger.info("Using NVIDIA APEX DistributedDataParallel.")
model = ApexDDP(model, delay_allreduce=True)
else:
if args.local_rank == 0:
_logger.info("Using native Torch DistributedDataParallel.")
model = NativeDDP(model, device_ids=[args.local_rank]) # can use device str in Torch >= 1.1
# NOTE: EMA model does not need to be wrapped by DDP
# setup mixup / cutmix
collate_fn = None
mixup_fn = None
mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
if mixup_active:
mixup_args = dict(
mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, cutmix_minmax=args.cutmix_minmax,
prob=args.mixup_prob, switch_prob=args.mixup_switch_prob, mode=args.mixup_mode,
label_smoothing=args.smoothing, num_classes=args.num_classes)
if args.prefetcher:
assert not num_aug_splits # collate conflict (need to support deinterleaving in collate mixup)
collate_fn = FastCollateMixup(**mixup_args)
else:
mixup_fn = Mixup(**mixup_args)
# create the train and eval datasets
train_dir = os.path.join(args.data, 'train')
if not os.path.exists(train_dir):
_logger.error('Training folder does not exist at: {}'.format(train_dir))
exit(1)
dataset_train = Dataset(train_dir)
eval_dir = os.path.join(args.data, 'val')
if not os.path.isdir(eval_dir):
eval_dir = os.path.join(args.data, 'validation')
if not os.path.isdir(eval_dir):
_logger.error('Validation folder does not exist at: {}'.format(eval_dir))
exit(1)
dataset_eval = Dataset(eval_dir)
# wrap dataset in AugMix helper
if num_aug_splits > 1:
dataset_train = AugMixDataset(dataset_train, num_splits=num_aug_splits)
# create data loaders w/ augmentation pipeline
train_interpolation = args.train_interpolation
if args.no_aug or not train_interpolation:
train_interpolation = data_config['interpolation']
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
no_aug=args.no_aug,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
re_split=args.resplit,
scale=args.scale,
ratio=args.ratio,
hflip=args.hflip,
vflip=args.vflip,
color_jitter=args.color_jitter,
auto_augment=args.aa,
num_aug_splits=num_aug_splits,
interpolation=train_interpolation,
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
pin_memory=args.pin_mem,
use_multi_epochs_loader=args.use_multi_epochs_loader
)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=args.validation_batch_size_multiplier * args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
crop_pct=data_config['crop_pct'],
pin_memory=args.pin_mem,
)
# setup learning rate schedule and starting epoch
lr_scheduler, num_epochs = create_scheduler(args, optimizer, dataset_train)
start_epoch = 0
if args.start_epoch is not None:
# a specified start_epoch will always override the resume epoch
start_epoch = args.start_epoch
elif resume_epoch is not None:
start_epoch = resume_epoch
if lr_scheduler is not None and start_epoch > 0:
lr_scheduler.step(start_epoch)
if cp_loaded is not None:
lr_scheduler.load_state_dict(cp_loaded['scheduler'])
if args.local_rank == 0:
_logger.info('Scheduled epochs: {}'.format(num_epochs))
# setup loss function
if args.jsd:
assert num_aug_splits > 1 # JSD only valid with aug splits set
train_loss_fn = JsdCrossEntropy(num_splits=num_aug_splits, smoothing=args.smoothing).cuda()
elif mixup_active:
# smoothing is handled with mixup target transform
train_loss_fn = SoftTargetCrossEntropy().cuda()
elif args.smoothing:
train_loss_fn = LabelSmoothingCrossEntropy(smoothing=args.smoothing).cuda()
else:
train_loss_fn = nn.CrossEntropyLoss().cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
# setup checkpoint saver and eval metric tracking
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
output_dir = ''
if args.local_rank == 0:
output_base = args.output if args.output else './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"),
args.model,
str(data_config['input_size'][-1])
])
output_dir = get_outdir(output_base, 'train', exp_name)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(
model=model, optimizer=optimizer, args=args, model_ema=model_ema, amp_scaler=loss_scaler,
checkpoint_dir=output_dir, recovery_dir=args.resume, decreasing=decreasing)
with open(os.path.join(output_dir, 'args.yaml'), 'w') as f:
f.write(args_text)
fieldnames = ['seed', 'weight_init', 'actfun', 'epoch', 'max_lr', 'lr', 'train_loss', 'eval_loss', 'eval_acc1', 'eval_acc5', 'ema']
filename = 'output'
if args.actfun != 'swish':
filename = '{}_'.format(args.actfun) + filename
outfile_path = os.path.join(args.output, filename) + '.csv'
if not os.path.exists(outfile_path):
with open(outfile_path, mode='w') as out_file:
writer = csv.DictWriter(out_file, fieldnames=fieldnames, lineterminator='\n')
writer.writeheader()
try:
for epoch in range(start_epoch, num_epochs):
if os.path.exists(args.check_path):
amp_loss = None
if use_amp == 'native':
amp_loss = loss_scaler.state_dict()
elif use_amp == 'apex':
amp_loss = amp.state_dict()
if model_ema is not None:
ema_save = model_ema.state_dict()
else:
ema_save = None
torch.save({'model': model_raw.state_dict(),
'model_ema': ema_save,
'optimizer': optimizer.state_dict(),
'scheduler': lr_scheduler.state_dict(),
'epoch': epoch,
'amp': amp_loss
}, check_path)
_logger.info('============ SAVED CHECKPOINT: Epoch {}'.format(epoch))
if args.distributed:
loader_train.sampler.set_epoch(epoch)
train_metrics = train_epoch(
epoch, model, loader_train, optimizer, train_loss_fn, args,
lr_scheduler=lr_scheduler, saver=saver, output_dir=output_dir,
amp_autocast=amp_autocast, loss_scaler=loss_scaler, model_ema=model_ema, mixup_fn=mixup_fn,
pre_model=pre_model)
if args.distributed and args.dist_bn in ('broadcast', 'reduce'):
if args.local_rank == 0:
_logger.info("Distributing BatchNorm running means and vars")
distribute_bn(model, args.world_size, args.dist_bn == 'reduce')
eval_metrics = validate(model, loader_eval, validate_loss_fn, args, amp_autocast=amp_autocast,
pre_model=pre_model)
with open(outfile_path, mode='a') as out_file:
writer = csv.DictWriter(out_file, fieldnames=fieldnames, lineterminator='\n')
writer.writerow({'seed': args.seed,
'actfun': args.actfun,
'epoch': epoch,
'lr': train_metrics['lr'],
'train_loss': train_metrics['loss'],
'eval_loss': eval_metrics['loss'],
'eval_acc1': eval_metrics['top1'],
'eval_acc5': eval_metrics['top5'],
'ema': False
})
if model_ema is not None and not args.model_ema_force_cpu:
if args.distributed and args.dist_bn in ('broadcast', 'reduce'):
distribute_bn(model_ema, args.world_size, args.dist_bn == 'reduce')
ema_eval_metrics = validate(
model_ema.module, loader_eval, validate_loss_fn, args, amp_autocast=amp_autocast, log_suffix=' (EMA)',
pre_model=pre_model)
eval_metrics = ema_eval_metrics
with open(outfile_path, mode='a') as out_file:
writer = csv.DictWriter(out_file, fieldnames=fieldnames, lineterminator='\n')
writer.writerow({'seed': args.seed,
'weight_init': args.weight_init,
'actfun': args.actfun,
'epoch': epoch,
'max_lr': args.lr,
'lr': train_metrics['lr'],
'train_loss': train_metrics['loss'],
'eval_loss': eval_metrics['loss'],
'eval_acc1': eval_metrics['top1'],
'eval_acc5': eval_metrics['top5'],
'ema': True
})
if lr_scheduler is not None and args.sched != 'onecycle':
# step LR for next epoch
lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
update_summary(
args.seed, epoch, args.lr, args.epochs, args.batch_size, args.actfun,
train_metrics, eval_metrics, os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
best_metric, best_epoch = saver.save_checkpoint(epoch, metric=save_metric)
except KeyboardInterrupt:
pass
if best_metric is not None:
_logger.info('*** Best metric: {0} (epoch {1})'.format(best_metric, best_epoch))
def _resume_fit(self, train_data, val_data, time_limit=math.inf):
tic = time.time()
# TODO: regression not implemented
if self._problem_type != REGRESSION and (not self.classes or not self.num_class):
raise ValueError('This is a classification problem and we are not able to determine classes of dataset')
if max(self.start_epoch, self.epoch) >= self.epochs:
return {'time': self._time_elapsed}
# wrap DP if possible
if self.found_gpu:
self.net = torch.nn.DataParallel(self.net, device_ids=[int(i) for i in self.valid_gpus])
self.net = self.net.to(self.ctx[0])
# prepare dataset
train_dataset = train_data.to_torch()
val_dataset = val_data.to_torch()
# setup mixup / cutmix
self._collate_fn = None
self._mixup_fn = None
self.mixup_active = self._augmentation_cfg.mixup > 0 or self._augmentation_cfg.cutmix > 0. or self._augmentation_cfg.cutmix_minmax is not None
if self.mixup_active:
mixup_args = dict(
mixup_alpha=self._augmentation_cfg.mixup, cutmix_alpha=self._augmentation_cfg.cutmix,
cutmix_minmax=self._augmentation_cfg.cutmix_minmax, prob=self._augmentation_cfg.mixup_prob,
switch_prob=self._augmentation_cfg.mixup_switch_prob, mode=self._augmentation_cfg.mixup_mode,
label_smoothing=self._augmentation_cfg.smoothing, num_classes=self.num_class)
if self._misc_cfg.prefetcher:
self._collate_fn = FastCollateMixup(**mixup_args)
else:
self._mixup_fn = Mixup(**mixup_args)
# create data loaders w/ augmentation pipeiine
train_interpolation = self._augmentation_cfg.train_interpolation
if self._augmentation_cfg.no_aug or not train_interpolation:
train_interpolation = self._data_cfg.interpolation
train_loader = create_loader(
train_dataset,
input_size=self._data_cfg.input_size,
batch_size=self._train_cfg.batch_size,
is_training=True,
use_prefetcher=self._misc_cfg.prefetcher,
no_aug=self._augmentation_cfg.no_aug,
scale=self._augmentation_cfg.scale,
ratio=self._augmentation_cfg.ratio,
hflip=self._augmentation_cfg.hflip,
vflip=self._augmentation_cfg.vflip,
color_jitter=self._augmentation_cfg.color_jitter,
auto_augment=self._augmentation_cfg.auto_augment,
interpolation=train_interpolation,
mean=self._data_cfg.mean,
std=self._data_cfg.std,
num_workers=self._misc_cfg.num_workers,
distributed=False,
collate_fn=self._collate_fn,
pin_memory=self._misc_cfg.pin_mem,
use_multi_epochs_loader=self._misc_cfg.use_multi_epochs_loader
)
val_loader = create_loader(
val_dataset,
input_size=self._data_cfg.input_size,
batch_size=self._data_cfg.validation_batch_size_multiplier * self._train_cfg.batch_size,
is_training=False,
use_prefetcher=self._misc_cfg.prefetcher,
interpolation=self._data_cfg.interpolation,
mean=self._data_cfg.mean,
std=self._data_cfg.std,
num_workers=self._misc_cfg.num_workers,
distributed=False,
crop_pct=self._data_cfg.crop_pct,
pin_memory=self._misc_cfg.pin_mem,
)
self._time_elapsed += time.time() - tic
return self._train_loop(train_loader, val_loader, time_limit=time_limit)
def main():
global args
setup_default_logging()
args, args_text = _parse_args()
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
if args.distributed and args.num_gpu > 1:
logging.warning(
'Using more than one GPU per process in distributed mode is not allowed. Setting num_gpu to 1.'
)
args.num_gpu = 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.num_gpu = 1
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
assert args.rank >= 0
if args.distributed:
logging.info(
'Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
logging.info('Training with a single process on %d GPUs.' %
args.num_gpu)
torch.manual_seed(args.seed + args.rank)
np.random.seed(args.seed + args.rank)
model = create_model(args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
drop_rate=args.drop,
global_pool=args.gp,
bn_tf=args.bn_tf,
bn_momentum=args.bn_momentum,
bn_eps=args.bn_eps,
checkpoint_path=args.initial_checkpoint)
if args.binarizable:
Model_binary_patch(model)
if args.local_rank == 0:
logging.info('Model %s created, param count: %d' %
(args.model, sum([m.numel()
for m in model.parameters()])))
data_config = resolve_data_config(vars(args),
model=model,
verbose=args.local_rank == 0)
if args.num_gpu > 1:
if args.amp:
logging.warning(
'AMP does not work well with nn.DataParallel, disabling. Use distributed mode for multi-GPU AMP.'
)
args.amp = False
model = nn.DataParallel(model,
device_ids=list(range(args.num_gpu))).cuda()
else:
model.cuda()
optimizer = create_optimizer(args, model)
use_amp = False
if has_apex and args.amp:
print('Using amp with --opt-level {}.'.format(args.opt_level))
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.opt_level)
use_amp = True
else:
print('Do NOT use amp.')
if args.local_rank == 0:
logging.info('NVIDIA APEX {}. AMP {}.'.format(
'installed' if has_apex else 'not installed',
'on' if use_amp else 'off'))
# optionally resume from a checkpoint
resume_state = {}
resume_epoch = None
if args.resume:
resume_state, resume_epoch = resume_checkpoint(model, args.resume)
if resume_state and not args.no_resume_opt:
if 'optimizer' in resume_state:
if args.local_rank == 0:
logging.info('Restoring Optimizer state from checkpoint')
optimizer.load_state_dict(resume_state['optimizer'])
if use_amp and 'amp' in resume_state and 'load_state_dict' in amp.__dict__:
if args.local_rank == 0:
logging.info('Restoring NVIDIA AMP state from checkpoint')
amp.load_state_dict(resume_state['amp'])
resume_state = None
if args.freeze_binary:
Model_freeze_binary(model)
if args.distributed:
if args.sync_bn:
try:
if has_apex:
model = convert_syncbn_model(model)
else:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
model)
if args.local_rank == 0:
logging.info(
'Converted model to use Synchronized BatchNorm.')
except Exception as e:
logging.error(
'Failed to enable Synchronized BatchNorm. Install Apex or Torch >= 1.1'
)
if has_apex:
model = DDP(model, delay_allreduce=True)
else:
if args.local_rank == 0:
logging.info(
"Using torch DistributedDataParallel. Install NVIDIA Apex for Apex DDP."
)
model = DDP(model,
device_ids=[args.local_rank
]) # can use device str in Torch >= 1.1
# NOTE: EMA model does not need to be wrapped by DDP
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
print(num_epochs)
# start_epoch = 0 #
if args.start_epoch is not None:
# a specified start_epoch will always override the resume epoch
start_epoch = args.start_epoch
elif resume_epoch is not None:
start_epoch = resume_epoch
if args.reset_lr_scheduler is not None:
lr_scheduler.base_values = len(
lr_scheduler.base_values) * [args.reset_lr_scheduler]
lr_scheduler.step(start_epoch)
if lr_scheduler is not None and start_epoch > 0:
lr_scheduler.step(start_epoch)
if args.local_rank == 0:
logging.info('Scheduled epochs: {}'.format(num_epochs))
# Using pruner to get sparse weights
if args.prune:
pruner = Pruner_mixed(model, 0, 100, args.pruner)
else:
pruner = None
dataset_train = torchvision.datasets.CIFAR100(root='~/Downloads/CIFAR100',
train=True,
download=True)
collate_fn = None
if args.prefetcher and args.mixup > 0:
collate_fn = FastCollateMixup(args.mixup, args.smoothing,
args.num_classes)
loader_train = create_loader_CIFAR100(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
rand_erase_prob=args.reprob,
rand_erase_mode=args.remode,
rand_erase_count=args.recount,
color_jitter=args.color_jitter,
auto_augment=args.aa,
interpolation='random',
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
is_clean_data=args.clean_train,
)
dataset_eval = torchvision.datasets.CIFAR100(root='~/Downloads/CIFAR100',
train=False,
download=True)
loader_eval = create_loader_CIFAR100(
dataset_eval,
input_size=data_config['input_size'],
batch_size=4 * args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
)
if args.mixup > 0.:
# smoothing is handled with mixup label transform
train_loss_fn = SoftTargetCrossEntropy(
multiplier=args.softmax_multiplier).cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
elif args.smoothing:
train_loss_fn = LabelSmoothingCrossEntropy(
smoothing=args.smoothing).cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
else:
train_loss_fn = nn.CrossEntropyLoss().cuda()
validate_loss_fn = train_loss_fn
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
saver_last_10_epochs = None
output_dir = ''
if args.local_rank == 0:
output_base = args.output if args.output else './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"), args.model,
str(data_config['input_size'][-1])
])
output_dir = get_outdir(output_base, 'train', exp_name)
decreasing = True if eval_metric == 'loss' else False
os.makedirs(output_dir + '/Top')
os.makedirs(output_dir + '/Last')
saver = CheckpointSaver(
checkpoint_dir=output_dir + '/Top',
decreasing=decreasing,
max_history=10) # Save the results of the top 10 epochs
saver_last_10_epochs = CheckpointSaver(
checkpoint_dir=output_dir + '/Last',
decreasing=decreasing,
max_history=10) # Save the results of the last 10 epochs
with open(os.path.join(output_dir, 'args.yaml'), 'w') as f:
f.write(args_text)
f.write('==============================\n')
f.write(model.__str__())
# if pruner:
# f.write('\n Sparsity \n')
# #f.write(pruner.threshold_dict.__str__())
# f.write('\n pruner.start_epoch={}, pruner.end_epoch={}'.format(pruner.start_epoch, pruner.end_epoch))
tensorboard_writer = SummaryWriter(output_dir)
try:
for epoch in range(start_epoch, num_epochs):
global alpha
alpha = get_alpha(epoch, args)
if args.distributed:
loader_train.sampler.set_epoch(epoch)
if pruner:
pruner.on_epoch_begin(epoch) # pruning
train_metrics = train_epoch(epoch,
model,
loader_train,
optimizer,
train_loss_fn,
args,
lr_scheduler=lr_scheduler,
saver=saver,
output_dir=output_dir,
use_amp=use_amp,
tensorboard_writer=tensorboard_writer,
pruner=pruner)
if pruner:
pruner.print_statistics()
eval_metrics = validate(model,
loader_eval,
validate_loss_fn,
args,
tensorboard_writer=tensorboard_writer,
epoch=epoch)
if lr_scheduler is not None:
# step LR for next epoch
lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
update_summary(epoch,
train_metrics,
eval_metrics,
os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
best_metric, best_epoch = saver.save_checkpoint(
model,
optimizer,
args,
epoch=epoch,
metric=save_metric,
use_amp=use_amp)
if saver_last_10_epochs is not None:
# save the checkpoint in the last 5 epochs
_, _ = saver_last_10_epochs.save_checkpoint(model,
optimizer,
args,
epoch=epoch,
metric=epoch,
use_amp=use_amp)
except KeyboardInterrupt:
pass
if best_metric is not None:
logging.info('*** Best metric: {0} (epoch {1})'.format(
best_metric, best_epoch))
def main(fold_i=0, data_=None, train_index=None, val_index=None):
setup_default_logging()
args, args_text = _parse_args()
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
best_score = 0.0
args.output = args.output + 'fold_' + str(fold_i)
if args.distributed:
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
if fold_i == 0:
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
_logger.info(
'Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
_logger.info('Training with a single process on 1 GPUs.')
assert args.rank >= 0
# resolve AMP arguments based on PyTorch / Apex availability
use_amp = None
if args.amp:
# for backwards compat, `--amp` arg tries apex before native amp
if has_apex:
args.apex_amp = True
elif has_native_amp:
args.native_amp = True
if args.apex_amp and has_apex:
use_amp = 'apex'
elif args.native_amp and has_native_amp:
use_amp = 'native'
elif args.apex_amp or args.native_amp:
_logger.warning(
"Neither APEX or native Torch AMP is available, using float32. "
"Install NVIDA apex or upgrade to PyTorch 1.6")
torch.manual_seed(args.seed + args.rank)
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,
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.local_rank == 0:
_logger.info('Model %s created, param count: %d' %
(args.model, sum([m.numel()
for m in model.parameters()])))
data_config = resolve_data_config(vars(args),
model=model,
verbose=args.local_rank == 0)
# setup augmentation batch splits for contrastive loss or split bn
num_aug_splits = 0
if args.aug_splits > 0:
assert args.aug_splits > 1, 'A split of 1 makes no sense'
num_aug_splits = args.aug_splits
# enable split bn (separate bn stats per batch-portion)
if args.split_bn:
assert num_aug_splits > 1 or args.resplit
model = convert_splitbn_model(model, max(num_aug_splits, 2))
# move model to GPU, enable channels last layout if set
model = nn.DataParallel(model)
model.cuda()
if args.channels_last:
model = model.to(memory_format=torch.channels_last)
# setup synchronized BatchNorm for distributed training
if args.distributed and args.sync_bn:
assert not args.split_bn
if has_apex and use_amp != 'native':
# Apex SyncBN preferred unless native amp is activated
model = convert_syncbn_model(model)
else:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
if args.local_rank == 0:
_logger.info(
'Converted model to use Synchronized BatchNorm. WARNING: You may have issues if using '
'zero initialized BN layers (enabled by default for ResNets) while sync-bn enabled.'
)
if args.torchscript:
assert not use_amp == 'apex', 'Cannot use APEX AMP with torchscripted model'
assert not args.sync_bn, 'Cannot use SyncBatchNorm with torchscripted model'
model = torch.jit.script(model)
optimizer = create_optimizer(args, model)
#optimizer = torch.optim.SGD(model.parameters(), lr=0.1, weight_decay=1e-6)
# setup automatic mixed-precision (AMP) loss scaling and op casting
amp_autocast = suppress # do nothing
loss_scaler = None
if use_amp == 'apex':
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
loss_scaler = ApexScaler()
if args.local_rank == 0:
_logger.info('Using NVIDIA APEX AMP. Training in mixed precision.')
elif use_amp == 'native':
amp_autocast = torch.cuda.amp.autocast
loss_scaler = NativeScaler()
if args.local_rank == 0:
_logger.info(
'Using native Torch AMP. Training in mixed precision.')
else:
if args.local_rank == 0:
_logger.info('AMP not enabled. Training in float32.')
# optionally resume from a checkpoint
resume_epoch = None
if args.resume:
resume_epoch = resume_checkpoint(
model,
args.resume,
optimizer=None if args.no_resume_opt else optimizer,
loss_scaler=None if args.no_resume_opt else loss_scaler,
log_info=args.local_rank == 0)
# setup exponential moving average of model weights, SWA could be used here too
model_ema = None
if args.model_ema:
# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
model_ema = ModelEmaV2(
model,
decay=args.model_ema_decay,
device='cpu' if args.model_ema_force_cpu else None)
if args.resume:
load_checkpoint(model_ema.module, args.resume, use_ema=True)
# setup distributed training
if args.distributed:
if has_apex and use_amp != 'native':
# Apex DDP preferred unless native amp is activated
if args.local_rank == 0:
_logger.info("Using NVIDIA APEX DistributedDataParallel.")
model = ApexDDP(model, delay_allreduce=True)
else:
if args.local_rank == 0:
_logger.info("Using native Torch DistributedDataParallel.")
model = NativeDDP(model, device_ids=[
args.local_rank
]) # can use device str in Torch >= 1.1
# NOTE: EMA model does not need to be wrapped by DDP
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
# lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=1, eta_min=1e-6, last_epoch=-1)
if args.local_rank == 0:
_logger.info('Scheduled epochs: {}'.format(20))
##create DataLoader
train_trans = get_riadd_train_transforms(args)
valid_trans = get_riadd_valid_transforms(args)
train_data = data_.iloc[train_index, :].reset_index(drop=True)
dataset_train = RiaddDataSet(image_ids=train_data, baseImgPath=args.data)
val_data = data_.iloc[val_index, :].reset_index(drop=True)
dataset_eval = RiaddDataSet(image_ids=val_data, baseImgPath=args.data)
# setup mixup / cutmix
collate_fn = None
mixup_fn = None
mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
if mixup_active:
mixup_args = dict(mixup_alpha=args.mixup,
cutmix_alpha=args.cutmix,
cutmix_minmax=args.cutmix_minmax,
prob=args.mixup_prob,
switch_prob=args.mixup_switch_prob,
mode=args.mixup_mode,
label_smoothing=args.smoothing,
num_classes=args.num_classes)
if args.prefetcher:
assert not num_aug_splits # collate conflict (need to support deinterleaving in collate mixup)
collate_fn = FastCollateMixup(**mixup_args)
else:
mixup_fn = Mixup(**mixup_args)
# wrap dataset in AugMix helper
if num_aug_splits > 1:
dataset_train = AugMixDataset(dataset_train, num_splits=num_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']
train_trans = get_riadd_train_transforms(args)
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
no_aug=args.no_aug,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
re_split=args.resplit,
scale=args.scale,
ratio=args.ratio,
hflip=args.hflip,
vflip=args.vflip,
color_jitter=args.color_jitter,
auto_augment=args.aa,
num_aug_splits=num_aug_splits,
interpolation=train_interpolation,
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
pin_memory=args.pin_mem,
use_multi_epochs_loader=args.use_multi_epochs_loader,
transform=train_trans)
valid_trans = get_riadd_valid_transforms(args)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=args.validation_batch_size_multiplier * args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
crop_pct=data_config['crop_pct'],
pin_memory=args.pin_mem,
transform=valid_trans)
# # setup loss function
# if args.jsd:
# assert num_aug_splits > 1 # JSD only valid with aug splits set
# train_loss_fn = JsdCrossEntropy(num_splits=num_aug_splits, smoothing=args.smoothing).cuda()
# elif mixup_active:
# # smoothing is handled with mixup target transform
# train_loss_fn = SoftTargetCrossEntropy().cuda()
# elif args.smoothing:
# train_loss_fn = LabelSmoothingCrossEntropy(smoothing=args.smoothing).cuda()
# else:
# train_loss_fn = nn.CrossEntropyLoss().cuda()
validate_loss_fn = nn.BCEWithLogitsLoss().cuda()
train_loss_fn = nn.BCEWithLogitsLoss().cuda()
# setup checkpoint saver and eval metric tracking
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
vis = None
output_dir = ''
if args.local_rank == 0:
output_base = args.output if args.output else './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"), args.model,
str(data_config['input_size'][-1])
])
output_dir = get_outdir(output_base, 'train', exp_name)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(model=model,
optimizer=optimizer,
args=args,
model_ema=model_ema,
amp_scaler=loss_scaler,
checkpoint_dir=output_dir,
recovery_dir=output_dir,
decreasing=decreasing)
with open(os.path.join(output_dir, 'args.yaml'), 'w') as f:
f.write(args_text)
vis = Visualizer(env=args.output)
try:
for epoch in range(0, args.epochs):
if args.distributed:
loader_train.sampler.set_epoch(epoch)
train_metrics = train_epoch(epoch,
model,
loader_train,
optimizer,
train_loss_fn,
args,
lr_scheduler=lr_scheduler,
saver=saver,
output_dir=output_dir,
amp_autocast=amp_autocast,
loss_scaler=loss_scaler,
model_ema=model_ema,
mixup_fn=mixup_fn)
if args.distributed and args.dist_bn in ('broadcast', 'reduce'):
if args.local_rank == 0:
_logger.info(
"Distributing BatchNorm running means and vars")
distribute_bn(model, args.world_size, args.dist_bn == 'reduce')
eval_metrics = validate(model,
loader_eval,
validate_loss_fn,
args,
amp_autocast=amp_autocast)
score, scores = get_score(eval_metrics['valid_label'],
eval_metrics['predictions'])
##visdom
if vis is not None:
vis.plot_curves({'None': epoch},
iters=epoch,
title='None',
xlabel='iters',
ylabel='None')
vis.plot_curves(
{'learing rate': optimizer.param_groups[0]['lr']},
iters=epoch,
title='lr',
xlabel='iters',
ylabel='learing rate')
vis.plot_curves({'train loss': float(train_metrics['loss'])},
iters=epoch,
title='train loss',
xlabel='iters',
ylabel='train loss')
vis.plot_curves({'val loss': float(eval_metrics['loss'])},
iters=epoch,
title='val loss',
xlabel='iters',
ylabel='val loss')
vis.plot_curves({'val score': float(score)},
iters=epoch,
title='val score',
xlabel='iters',
ylabel='val score')
if model_ema is not None and not args.model_ema_force_cpu:
if args.distributed and args.dist_bn in ('broadcast',
'reduce'):
distribute_bn(model_ema, args.world_size,
args.dist_bn == 'reduce')
ema_eval_metrics = validate(model_ema.module,
loader_eval,
validate_loss_fn,
args,
amp_autocast=amp_autocast,
log_suffix=' (EMA)')
eval_metrics = ema_eval_metrics
if lr_scheduler is not None:
# step LR for next epoch
# lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
lr_scheduler.step(epoch + 1, score)
update_summary(epoch,
train_metrics,
eval_metrics,
os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None)
if saver is not None and score > best_score:
# save proper checkpoint with eval metric
best_score = score
save_metric = best_score
best_metric, best_epoch = saver.save_checkpoint(
epoch, metric=save_metric)
del model
del optimizer
torch.cuda.empty_cache()
except KeyboardInterrupt:
pass
if best_metric is not None:
_logger.info('*** Best metric: {0} (epoch {1})'.format(
best_metric, best_epoch))
def main():
setup_default_logging()
args, args_text = _parse_args()
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
if args.distributed and args.num_gpu > 1:
_logger.warning(
'Using more than one GPU per process in distributed mode is not allowed.Setting num_gpu to 1.')
args.num_gpu = 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.num_gpu = 1
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
assert args.rank >= 0
if args.distributed:
_logger.info('Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
_logger.info('Training with a single process on %d GPUs.' % args.num_gpu)
torch.manual_seed(args.seed + args.rank)
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,
global_pool=args.gp,
bn_tf=args.bn_tf,
bn_momentum=args.bn_momentum,
bn_eps=args.bn_eps,
checkpoint_path=args.initial_checkpoint)
if args.local_rank == 0:
_logger.info('Model %s created, param count: %d' %
(args.model, sum([m.numel() for m in model.parameters()])))
data_config = resolve_data_config(vars(args), model=model, verbose=args.local_rank == 0)
num_aug_splits = 0
if args.aug_splits > 0:
assert args.aug_splits > 1, 'A split of 1 makes no sense'
num_aug_splits = args.aug_splits
if args.split_bn:
assert num_aug_splits > 1 or args.resplit
model = convert_splitbn_model(model, max(num_aug_splits, 2))
use_amp = None
if args.amp:
# for backwards compat, `--amp` arg tries apex before native amp
if has_apex:
args.apex_amp = True
elif has_native_amp:
args.native_amp = True
if args.apex_amp and has_apex:
use_amp = 'apex'
elif args.native_amp and has_native_amp:
use_amp = 'native'
elif args.apex_amp or args.native_amp:
_logger.warning("Neither APEX or native Torch AMP is available, using float32. "
"Install NVIDA apex or upgrade to PyTorch 1.6")
if args.num_gpu > 1:
if use_amp == 'apex':
_logger.warning(
'Apex AMP does not work well with nn.DataParallel, disabling. Use DDP or Torch AMP.')
use_amp = None
model = nn.DataParallel(model, device_ids=list(range(args.num_gpu))).cuda()
assert not args.channels_last, "Channels last not supported with DP, use DDP."
else:
model.cuda()
if args.channels_last:
model = model.to(memory_format=torch.channels_last)
optimizer = create_optimizer(args, model)
amp_autocast = suppress # do nothing
loss_scaler = None
if use_amp == 'apex':
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
loss_scaler = ApexScaler()
if args.local_rank == 0:
_logger.info('Using NVIDIA APEX AMP. Training in mixed precision.')
elif use_amp == 'native':
amp_autocast = torch.cuda.amp.autocast
loss_scaler = NativeScaler()
if args.local_rank == 0:
_logger.info('Using native Torch AMP. Training in mixed precision.')
else:
if args.local_rank == 0:
_logger.info('AMP not enabled. Training in float32.')
# optionally resume from a checkpoint
resume_epoch = None
if args.resume:
resume_epoch = resume_checkpoint(
model, args.resume,
optimizer=None if args.no_resume_opt else optimizer,
loss_scaler=None if args.no_resume_opt else loss_scaler,
log_info=args.local_rank == 0)
model_ema = None
if args.model_ema:
# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
model_ema = ModelEma(
model,
decay=args.model_ema_decay,
device='cpu' if args.model_ema_force_cpu else '',
resume=args.resume)
if args.distributed:
if args.sync_bn:
assert not args.split_bn
try:
if has_apex and use_amp != 'native':
# Apex SyncBN preferred unless native amp is activated
model = convert_syncbn_model(model)
else:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
if args.local_rank == 0:
_logger.info(
'Converted model to use Synchronized BatchNorm. WARNING: You may have issues if using '
'zero initialized BN layers (enabled by default for ResNets) while sync-bn enabled.')
except Exception as e:
_logger.error('Failed to enable Synchronized BatchNorm. Install Apex or Torch >= 1.1')
if has_apex and use_amp != 'native':
# Apex DDP preferred unless native amp is activated
if args.local_rank == 0:
_logger.info("Using NVIDIA APEX DistributedDataParallel.")
model = ApexDDP(model, delay_allreduce=True)
else:
if args.local_rank == 0:
_logger.info("Using native Torch DistributedDataParallel.")
model = NativeDDP(model, device_ids=[args.local_rank]) # can use device str in Torch >= 1.1
# NOTE: EMA model does not need to be wrapped by DDP
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
start_epoch = 0
if args.start_epoch is not None:
# a specified start_epoch will always override the resume epoch
start_epoch = args.start_epoch
elif resume_epoch is not None:
start_epoch = resume_epoch
if lr_scheduler is not None and start_epoch > 0:
lr_scheduler.step(start_epoch)
if args.local_rank == 0:
_logger.info('Scheduled epochs: {}'.format(num_epochs))
train_dir = os.path.join(args.data, 'train')
if not os.path.exists(train_dir):
_logger.error('Training folder does not exist at: {}'.format(train_dir))
exit(1)
dataset_train = Dataset(train_dir)
collate_fn = None
mixup_fn = None
mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
if mixup_active:
mixup_args = dict(
mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, cutmix_minmax=args.cutmix_minmax,
prob=args.mixup_prob, switch_prob=args.mixup_switch_prob, mode=args.mixup_mode,
label_smoothing=args.smoothing, num_classes=args.num_classes)
if args.prefetcher:
assert not num_aug_splits # collate conflict (need to support deinterleaving in collate mixup)
collate_fn = FastCollateMixup(**mixup_args)
else:
mixup_fn = Mixup(**mixup_args)
if num_aug_splits > 1:
dataset_train = AugMixDataset(dataset_train, num_splits=num_aug_splits)
train_interpolation = args.train_interpolation
if args.no_aug or not train_interpolation:
train_interpolation = data_config['interpolation']
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
no_aug=args.no_aug,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
re_split=args.resplit,
scale=args.scale,
ratio=args.ratio,
hflip=args.hflip,
vflip=args.vflip,
color_jitter=args.color_jitter,
auto_augment=args.aa,
num_aug_splits=num_aug_splits,
interpolation=train_interpolation,
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
pin_memory=args.pin_mem,
use_multi_epochs_loader=args.use_multi_epochs_loader
)
eval_dir = os.path.join(args.data, 'val')
if not os.path.isdir(eval_dir):
eval_dir = os.path.join(args.data, 'validation')
if not os.path.isdir(eval_dir):
_logger.error('Validation folder does not exist at: {}'.format(eval_dir))
exit(1)
dataset_eval = Dataset(eval_dir)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=args.validation_batch_size_multiplier * args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
crop_pct=data_config['crop_pct'],
pin_memory=args.pin_mem,
)
if args.jsd:
assert num_aug_splits > 1 # JSD only valid with aug splits set
train_loss_fn = JsdCrossEntropy(num_splits=num_aug_splits, smoothing=args.smoothing).cuda()
elif mixup_active:
# smoothing is handled with mixup target transform
train_loss_fn = SoftTargetCrossEntropy().cuda()
elif args.smoothing:
train_loss_fn = LabelSmoothingCrossEntropy(smoothing=args.smoothing).cuda()
else:
train_loss_fn = nn.CrossEntropyLoss().cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
output_dir = ''
if args.local_rank == 0:
output_base = args.output if args.output else './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"),
args.model,
str(data_config['input_size'][-1])
])
output_dir = get_outdir(output_base, 'train', exp_name)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(
model=model, optimizer=optimizer, args=args, model_ema=model_ema, amp_scaler=loss_scaler,
checkpoint_dir=output_dir, recovery_dir=output_dir, decreasing=decreasing)
with open(os.path.join(output_dir, 'args.yaml'), 'w') as f:
f.write(args_text)
try:
for epoch in range(start_epoch, num_epochs):
if args.distributed:
loader_train.sampler.set_epoch(epoch)
train_metrics = train_epoch(
epoch, model, loader_train, optimizer, train_loss_fn, args,
lr_scheduler=lr_scheduler, saver=saver, output_dir=output_dir,
amp_autocast=amp_autocast, loss_scaler=loss_scaler, model_ema=model_ema, mixup_fn=mixup_fn)
if args.distributed and args.dist_bn in ('broadcast', 'reduce'):
if args.local_rank == 0:
_logger.info("Distributing BatchNorm running means and vars")
distribute_bn(model, args.world_size, args.dist_bn == 'reduce')
eval_metrics = validate(model, loader_eval, validate_loss_fn, args, amp_autocast=amp_autocast)
if model_ema is not None and not args.model_ema_force_cpu:
if args.distributed and args.dist_bn in ('broadcast', 'reduce'):
distribute_bn(model_ema, args.world_size, args.dist_bn == 'reduce')
ema_eval_metrics = validate(
model_ema.ema, loader_eval, validate_loss_fn, args, amp_autocast=amp_autocast, log_suffix=' (EMA)')
eval_metrics = ema_eval_metrics
if lr_scheduler is not None:
# step LR for next epoch
lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
update_summary(
epoch, train_metrics, eval_metrics, os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
best_metric, best_epoch = saver.save_checkpoint(epoch, metric=save_metric)
# if saver.cmp(best_metric, save_metric):
# _logger.info(f"Metric is no longer improving [BEST: {best_metric}, CURRENT: {save_metric}]"
# f"\nFinishing training process")
# if epoch > 15:
# break
except KeyboardInterrupt:
pass
if best_metric is not None:
message = '*** Best metric: <{0:.2f}>, epoch: <{1}>, path: <{2}> ***'\
.format(best_metric, best_epoch, output_dir)
_logger.info(message)
print(message)
def train_imagenet():
torch.manual_seed(42)
device = xm.xla_device()
# model = get_model_property('model_fn')().to(device)
model = create_model(
FLAGS.model,
pretrained=FLAGS.pretrained,
num_classes=FLAGS.num_classes,
drop_rate=FLAGS.drop,
global_pool=FLAGS.gp,
bn_tf=FLAGS.bn_tf,
bn_momentum=FLAGS.bn_momentum,
bn_eps=FLAGS.bn_eps,
drop_connect_rate=0.2,
checkpoint_path=FLAGS.initial_checkpoint,
args = FLAGS).to(device)
model_ema=None
if FLAGS.model_ema:
# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
# import pdb; pdb.set_trace()
model_e = create_model(
FLAGS.model,
pretrained=FLAGS.pretrained,
num_classes=FLAGS.num_classes,
drop_rate=FLAGS.drop,
global_pool=FLAGS.gp,
bn_tf=FLAGS.bn_tf,
bn_momentum=FLAGS.bn_momentum,
bn_eps=FLAGS.bn_eps,
drop_connect_rate=0.2,
checkpoint_path=FLAGS.initial_checkpoint,
args = FLAGS).to(device)
model_ema = ModelEma(
model_e,
decay=FLAGS.model_ema_decay,
device='cpu' if FLAGS.model_ema_force_cpu else '',
resume=FLAGS.resume)
print('==> Preparing data..')
img_dim = 224
if FLAGS.fake_data:
train_dataset_len = 1200000 # Roughly the size of Imagenet dataset.
train_loader = xu.SampleGenerator(
data=(torch.zeros(FLAGS.batch_size, 3, img_dim, img_dim),
torch.zeros(FLAGS.batch_size, dtype=torch.int64)),
sample_count=train_dataset_len // FLAGS.batch_size //
xm.xrt_world_size())
test_loader = xu.SampleGenerator(
data=(torch.zeros(FLAGS.batch_size, 3, img_dim, img_dim),
torch.zeros(FLAGS.batch_size, dtype=torch.int64)),
sample_count=50000 // FLAGS.batch_size // xm.xrt_world_size())
# else:
# normalize = transforms.Normalize(
# mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
# train_dataset = torchvision.datasets.ImageFolder(
# os.path.join(FLAGS.data, 'train'),
# transforms.Compose([
# transforms.RandomResizedCrop(img_dim),
# transforms.RandomHorizontalFlip(),
# transforms.ToTensor(),
# normalize,
# ]))
# train_dataset_len = len(train_dataset.imgs)
# resize_dim = max(img_dim, 256)
# test_dataset = torchvision.datasets.ImageFolder(
# os.path.join(FLAGS.data, 'val'),
# # Matches Torchvision's eval transforms except Torchvision uses size
# # 256 resize for all models both here and in the train loader. Their
# # version crashes during training on 299x299 images, e.g. inception.
# transforms.Compose([
# transforms.Resize(resize_dim),
# transforms.CenterCrop(img_dim),
# transforms.ToTensor(),
# normalize,
# ]))
# train_sampler = None
# if xm.xrt_world_size() > 1:
# train_sampler = torch.utils.data.distributed.DistributedSampler(
# train_dataset,
# num_replicas=xm.xrt_world_size(),
# rank=xm.get_ordinal(),
# shuffle=True)
# train_loader = torch.utils.data.DataLoader(
# train_dataset,
# batch_size=FLAGS.batch_size,
# sampler=train_sampler,
# shuffle=False if train_sampler else True,
# num_workers=FLAGS.workers)
# test_loader = torch.utils.data.DataLoader(
# test_dataset,
# batch_size=FLAGS.batch_size,
# shuffle=False,
# num_workers=FLAGS.workers)
else:
train_dir = os.path.join(FLAGS.data, 'train')
data_config = resolve_data_config(model, FLAGS, verbose=FLAGS.local_rank == 0)
dataset_train = Dataset(train_dir)
collate_fn = None
if not FLAGS.no_prefetcher and FLAGS.mixup > 0:
collate_fn = FastCollateMixup(FLAGS.mixup, FLAGS.smoothing, FLAGS.num_classes)
train_loader = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=FLAGS.batch_size,
is_training=True,
use_prefetcher=not FLAGS.no_prefetcher,
rand_erase_prob=FLAGS.reprob,
rand_erase_mode=FLAGS.remode,
interpolation='bicubic', # FIXME cleanly resolve this? data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=FLAGS.workers,
distributed=FLAGS.distributed,
collate_fn=collate_fn,
use_auto_aug=FLAGS.auto_augment,
use_mixcut=FLAGS.mixcut,
)
eval_dir = os.path.join(FLAGS.data, 'val')
train_dataset_len = len(train_loader)
if not os.path.isdir(eval_dir):
logging.error('Validation folder does not exist at: {}'.format(eval_dir))
exit(1)
dataset_eval = Dataset(eval_dir)
test_loader = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size = FLAGS.batch_size,
is_training=False,
use_prefetcher=FLAGS.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=FLAGS.workers,
distributed=FLAGS.distributed,
)
writer = None
start_epoch = 0
if FLAGS.output and xm.is_master_ordinal():
writer = SummaryWriter(log_dir=FLAGS.output)
optimizer = create_optimizer(flags, model)
lr_scheduler, num_epochs = create_scheduler(flags, optimizer)
if start_epoch > 0:
lr_scheduler.step(start_epoch)
# optimizer = optim.SGD(
# model.parameters(),
# lr=FLAGS.lr,
# momentum=FLAGS.momentum,
# weight_decay=5e-4)
num_training_steps_per_epoch = train_dataset_len // (
FLAGS.batch_size * xm.xrt_world_size())
lr_scheduler = schedulers.wrap_optimizer_with_scheduler(
optimizer,
scheduler_type=getattr(FLAGS, 'lr_scheduler_type', None),
scheduler_divisor=getattr(FLAGS, 'lr_scheduler_divisor', None),
scheduler_divide_every_n_epochs=getattr(
FLAGS, 'lr_scheduler_divide_every_n_epochs', None),
num_steps_per_epoch=num_training_steps_per_epoch,
summary_writer=writer)
train_loss_fn = LabelSmoothingCrossEntropy(smoothing=flags.smoothing)
validate_loss_fn = nn.CrossEntropyLoss()
# loss_fn = nn.CrossEntropyLoss()
def train_loop_fn(loader):
tracker = xm.RateTracker()
model.train()
for x, (data, target) in loader:
optimizer.zero_grad()
output = model(data)
loss = train_loss_fn(output, target)
loss.backward()
xm.optimizer_step(optimizer)
tracker.add(FLAGS.batch_size)
if model_ema is not None:
model_ema.update(model)
if lr_scheduler:
lr_scheduler.step()
if x % FLAGS.log_steps == 0:
test_utils.print_training_update(device, x, loss.item(), tracker.rate(),
tracker.global_rate())
def test_loop_fn(loader):
total_samples = 0
correct = 0
model.eval()
for x, (data, target) in loader:
output = model(data)
pred = output.max(1, keepdim=True)[1]
correct += pred.eq(target.view_as(pred)).sum().item()
total_samples += data.size()[0]
accuracy = 100.0 * correct / total_samples
test_utils.print_test_update(device, accuracy)
return accuracy
def test_loop_fn_ema(loader):
total_samples = 0
correct = 0
model_ema.eval()
for x, (data, target) in loader:
output = model_ema(data)
pred = output.max(1, keepdim=True)[1]
correct += pred.eq(target.view_as(pred)).sum().item()
total_samples += data.size()[0]
accuracy = 100.0 * correct / total_samples
test_utils.print_test_update(device, accuracy)
return accuracy
accuracy = 0.0
for epoch in range(1, FLAGS.epochs + 1):
para_loader = dp.ParallelLoader(train_loader, [device])
train_loop_fn(para_loader.per_device_loader(device))
para_loader = dp.ParallelLoader(test_loader, [device])
accuracy = test_loop_fn(para_loader.per_device_loader(device))
print('Epoch: {}, Mean Accuracy: {:.2f}%'.format(epoch, accuracy))
if model_ema is not None:
accuracy = test_loop_fn_ema(para_loader.per_device_loader(device))
print('Epoch: {}, Mean Accuracy: {:.2f}%'.format(epoch, accuracy))
test_utils.add_scalar_to_summary(writer, 'Accuracy/test', accuracy, epoch)
if FLAGS.metrics_debug:
print(torch_xla._XLAC._xla_metrics_report())
return accuracy
def main():
setup_default_logging()
args, args_text = _parse_args()
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
if args.distributed and args.num_gpu > 1:
logging.warning(
'Using more than one GPU per process in distributed mode is not allowed. Setting num_gpu to 1.'
)
args.num_gpu = 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.num_gpu = 1
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
assert args.rank >= 0
if args.distributed:
logging.info(
'Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
logging.info('Training with a single process on %d GPUs.' %
args.num_gpu)
torch.manual_seed(args.seed + args.rank)
model = create_model(args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
drop_rate=args.drop,
drop_connect_rate=args.drop_connect,
global_pool=args.gp,
bn_tf=args.bn_tf,
bn_momentum=args.bn_momentum,
bn_eps=args.bn_eps,
checkpoint_path=args.initial_checkpoint)
if args.local_rank == 0:
logging.info('Model %s created, param count: %d' %
(args.model, sum([m.numel()
for m in model.parameters()])))
data_config = resolve_data_config(vars(args),
model=model,
verbose=args.local_rank == 0)
num_aug_splits = 0
if args.aug_splits > 0:
assert args.aug_splits > 1, 'A split of 1 makes no sense'
num_aug_splits = args.aug_splits
if args.split_bn:
assert num_aug_splits > 1 or args.resplit
model = convert_splitbn_model(model, max(num_aug_splits, 2))
if args.num_gpu > 1:
if args.amp:
logging.warning(
'AMP does not work well with nn.DataParallel, disabling. Use distributed mode for multi-GPU AMP.'
)
args.amp = False
model = nn.DataParallel(model,
device_ids=list(range(args.num_gpu))).cuda()
else:
model.cuda()
optimizer = create_optimizer(args, model)
use_amp = False
if has_apex and args.amp:
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
use_amp = True
if args.local_rank == 0:
logging.info('NVIDIA APEX {}. AMP {}.'.format(
'installed' if has_apex else 'not installed',
'on' if use_amp else 'off'))
# optionally resume from a checkpoint
resume_state = {}
resume_epoch = None
if args.resume:
resume_state, resume_epoch = resume_checkpoint(model, args.resume)
if resume_state and not args.no_resume_opt:
if 'optimizer' in resume_state:
if args.local_rank == 0:
logging.info('Restoring Optimizer state from checkpoint')
optimizer.load_state_dict(resume_state['optimizer'])
if use_amp and 'amp' in resume_state and 'load_state_dict' in amp.__dict__:
if args.local_rank == 0:
logging.info('Restoring NVIDIA AMP state from checkpoint')
amp.load_state_dict(resume_state['amp'])
del resume_state
model_ema = None
if args.model_ema:
# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
model_ema = ModelEma(model,
decay=args.model_ema_decay,
device='cpu' if args.model_ema_force_cpu else '',
resume=args.resume)
if args.distributed:
if args.sync_bn:
assert not args.split_bn
try:
if has_apex:
model = convert_syncbn_model(model)
else:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
model)
if args.local_rank == 0:
logging.info(
'Converted model to use Synchronized BatchNorm. WARNING: You may have issues if using '
'zero initialized BN layers (enabled by default for ResNets) while sync-bn enabled.'
)
except Exception as e:
logging.error(
'Failed to enable Synchronized BatchNorm. Install Apex or Torch >= 1.1'
)
if has_apex:
model = DDP(model, delay_allreduce=True)
else:
if args.local_rank == 0:
logging.info(
"Using torch DistributedDataParallel. Install NVIDIA Apex for Apex DDP."
)
model = DDP(model,
device_ids=[args.local_rank
]) # can use device str in Torch >= 1.1
# NOTE: EMA model does not need to be wrapped by DDP
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
start_epoch = 0
if args.start_epoch is not None:
# a specified start_epoch will always override the resume epoch
start_epoch = args.start_epoch
elif resume_epoch is not None:
start_epoch = resume_epoch
if lr_scheduler is not None and start_epoch > 0:
lr_scheduler.step(start_epoch)
if args.local_rank == 0:
logging.info('Scheduled epochs: {}'.format(num_epochs))
train_dir = os.path.join(args.data, 'train')
if not os.path.exists(train_dir):
logging.error(
'Training folder does not exist at: {}'.format(train_dir))
exit(1)
dataset_train = Dataset(train_dir)
collate_fn = None
if args.prefetcher and args.mixup > 0:
assert not num_aug_splits # collate conflict (need to support deinterleaving in collate mixup)
collate_fn = FastCollateMixup(args.mixup, args.smoothing,
args.num_classes)
if num_aug_splits > 1:
dataset_train = AugMixDataset(dataset_train, num_splits=num_aug_splits)
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
re_split=args.resplit,
color_jitter=args.color_jitter,
auto_augment=args.aa,
num_aug_splits=num_aug_splits,
interpolation=args.train_interpolation,
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
pin_memory=args.pin_mem,
)
eval_dir = os.path.join(args.data, 'val')
if not os.path.isdir(eval_dir):
eval_dir = os.path.join(args.data, 'validation')
if not os.path.isdir(eval_dir):
logging.error(
'Validation folder does not exist at: {}'.format(eval_dir))
exit(1)
dataset_eval = Dataset(eval_dir)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=4 * args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
crop_pct=data_config['crop_pct'],
pin_memory=args.pin_mem,
)
if args.jsd:
assert num_aug_splits > 1 # JSD only valid with aug splits set
train_loss_fn = JsdCrossEntropy(num_splits=num_aug_splits,
smoothing=args.smoothing).cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
elif args.mixup > 0.:
# smoothing is handled with mixup label transform
train_loss_fn = SoftTargetCrossEntropy().cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
elif args.smoothing:
train_loss_fn = LabelSmoothingCrossEntropy(
smoothing=args.smoothing).cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
else:
train_loss_fn = nn.CrossEntropyLoss().cuda()
validate_loss_fn = train_loss_fn
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
output_dir = ''
if args.local_rank == 0:
output_base = args.output if args.output else './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"), args.model,
str(data_config['input_size'][-1])
])
output_dir = get_outdir(output_base, 'train', exp_name)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(checkpoint_dir=output_dir,
decreasing=decreasing)
with open(os.path.join(output_dir, 'args.yaml'), 'w') as f:
f.write(args_text)
try:
for epoch in range(start_epoch, num_epochs):
if args.distributed:
loader_train.sampler.set_epoch(epoch)
train_metrics = train_epoch(epoch,
model,
loader_train,
optimizer,
train_loss_fn,
args,
lr_scheduler=lr_scheduler,
saver=saver,
output_dir=output_dir,
use_amp=use_amp,
model_ema=model_ema)
if args.distributed and args.dist_bn in ('broadcast', 'reduce'):
if args.local_rank == 0:
logging.info(
"Distributing BatchNorm running means and vars")
distribute_bn(model, args.world_size, args.dist_bn == 'reduce')
eval_metrics = validate(model, loader_eval, validate_loss_fn, args)
if model_ema is not None and not args.model_ema_force_cpu:
if args.distributed and args.dist_bn in ('broadcast',
'reduce'):
distribute_bn(model_ema, args.world_size,
args.dist_bn == 'reduce')
ema_eval_metrics = validate(model_ema.ema,
loader_eval,
validate_loss_fn,
args,
log_suffix=' (EMA)')
eval_metrics = ema_eval_metrics
if lr_scheduler is not None:
# step LR for next epoch
lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
update_summary(epoch,
train_metrics,
eval_metrics,
os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
best_metric, best_epoch = saver.save_checkpoint(
model,
optimizer,
args,
epoch=epoch,
model_ema=model_ema,
metric=save_metric,
use_amp=use_amp)
except KeyboardInterrupt:
pass
if best_metric is not None:
logging.info('*** Best metric: {0} (epoch {1})'.format(
best_metric, best_epoch))
def main():
setup_default_logging()
args, args_text = _parse_args()
args.prefetcher = not args.no_prefetcher
args.distributed = False
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
_logger.info('====================\n\n'
'Actfun: {}\n'
'LR: {}\n'
'Epochs: {}\n'
'p: {}\n'
'k: {}\n'
'g: {}\n'
'Extra channel multiplier: {}\n'
'Weight Init: {}\n'
'\n===================='.format(args.actfun, args.lr,
args.epochs, args.p, args.k,
args.g,
args.extra_channel_mult,
args.weight_init))
# ================================================================================= Loading models
pre_model = create_model(
args.model,
pretrained=True,
actfun='swish',
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,
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,
p=args.p,
k=args.k,
g=args.g,
extra_channel_mult=args.extra_channel_mult,
weight_init_name=args.weight_init,
partial_ho_actfun=args.partial_ho_actfun)
pre_model_layers = list(pre_model.children())
pre_model = torch.nn.Sequential(*pre_model_layers[:-1])
pre_model.to(device)
model = MLP.MLP(actfun=args.actfun,
input_dim=1280,
output_dim=args.num_classes,
k=args.k,
p=args.p,
g=args.g,
num_params=1_000_000,
permute_type='shuffle')
model.to(device)
# ================================================================================= Loading dataset
util.seed_all(args.seed)
if args.data == 'caltech101' and not os.path.exists('caltech101'):
dir_root = r'101_ObjectCategories'
dir_new = r'caltech101'
dir_new_train = os.path.join(dir_new, 'train')
dir_new_val = os.path.join(dir_new, 'val')
dir_new_test = os.path.join(dir_new, 'test')
if not os.path.exists(dir_new):
os.mkdir(dir_new)
os.mkdir(dir_new_train)
os.mkdir(dir_new_val)
os.mkdir(dir_new_test)
for dir2 in os.listdir(dir_root):
if dir2 != 'BACKGROUND_Google':
curr_path = os.path.join(dir_root, dir2)
new_path_train = os.path.join(dir_new_train, dir2)
new_path_val = os.path.join(dir_new_val, dir2)
new_path_test = os.path.join(dir_new_test, dir2)
if not os.path.exists(new_path_train):
os.mkdir(new_path_train)
if not os.path.exists(new_path_val):
os.mkdir(new_path_val)
if not os.path.exists(new_path_test):
os.mkdir(new_path_test)
train_upper = int(0.8 * len(os.listdir(curr_path)))
val_upper = int(0.9 * len(os.listdir(curr_path)))
curr_files_all = os.listdir(curr_path)
curr_files_train = curr_files_all[:train_upper]
curr_files_val = curr_files_all[train_upper:val_upper]
curr_files_test = curr_files_all[val_upper:]
for file in curr_files_train:
copyfile(os.path.join(curr_path, file),
os.path.join(new_path_train, file))
for file in curr_files_val:
copyfile(os.path.join(curr_path, file),
os.path.join(new_path_val, file))
for file in curr_files_test:
copyfile(os.path.join(curr_path, file),
os.path.join(new_path_test, file))
time.sleep(5)
# create the train and eval datasets
train_dir = os.path.join(args.data, 'train')
if not os.path.exists(train_dir):
_logger.error(
'Training folder does not exist at: {}'.format(train_dir))
exit(1)
dataset_train = Dataset(train_dir)
eval_dir = os.path.join(args.data, 'val')
if not os.path.isdir(eval_dir):
eval_dir = os.path.join(args.data, 'validation')
if not os.path.isdir(eval_dir):
_logger.error(
'Validation folder does not exist at: {}'.format(eval_dir))
exit(1)
dataset_eval = Dataset(eval_dir)
# setup augmentation batch splits for contrastive loss or split bn
num_aug_splits = 0
if args.aug_splits > 0:
assert args.aug_splits > 1, 'A split of 1 makes no sense'
num_aug_splits = args.aug_splits
# enable split bn (separate bn stats per batch-portion)
if args.split_bn:
assert num_aug_splits > 1 or args.resplit
model = convert_splitbn_model(model, max(num_aug_splits, 2))
# setup mixup / cutmix
collate_fn = None
mixup_fn = None
mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
if mixup_active:
mixup_args = dict(mixup_alpha=args.mixup,
cutmix_alpha=args.cutmix,
cutmix_minmax=args.cutmix_minmax,
prob=args.mixup_prob,
switch_prob=args.mixup_switch_prob,
mode=args.mixup_mode,
label_smoothing=args.smoothing,
num_classes=args.num_classes)
if args.prefetcher:
assert not num_aug_splits # collate conflict (need to support deinterleaving in collate mixup)
collate_fn = FastCollateMixup(**mixup_args)
else:
mixup_fn = Mixup(**mixup_args)
# create data loaders w/ augmentation pipeline
train_interpolation = args.train_interpolation
data_config = resolve_data_config(vars(args),
model=model,
verbose=args.local_rank == 0)
if args.no_aug or not train_interpolation:
train_interpolation = data_config['interpolation']
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
no_aug=args.no_aug,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
re_split=args.resplit,
scale=args.scale,
ratio=args.ratio,
hflip=args.hflip,
vflip=args.vflip,
color_jitter=args.color_jitter,
auto_augment=args.aa,
num_aug_splits=num_aug_splits,
interpolation=train_interpolation,
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
pin_memory=args.pin_mem,
use_multi_epochs_loader=args.use_multi_epochs_loader)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=args.validation_batch_size_multiplier * args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
crop_pct=data_config['crop_pct'],
pin_memory=args.pin_mem,
)
# ================================================================================= Optimizer / scheduler
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), weight_decay=1e-5)
scheduler = OneCycleLR(
optimizer,
max_lr=args.lr,
epochs=args.epochs,
steps_per_epoch=int(math.floor(len(dataset_train) / args.batch_size)),
cycle_momentum=False)
# ================================================================================= Save file / checkpoints
fieldnames = [
'dataset', 'seed', 'epoch', 'time', 'actfun', 'model', 'batch_size',
'alpha_primes', 'alphas', 'num_params', 'k', 'p', 'g', 'perm_method',
'gen_gap', 'epoch_train_loss', 'epoch_train_acc',
'epoch_aug_train_loss', 'epoch_aug_train_acc', 'epoch_val_loss',
'epoch_val_acc', 'curr_lr', 'found_lr', 'epochs'
]
filename = 'out_{}_{}_{}_{}'.format(datetime.date.today(), args.actfun,
args.data, args.seed)
outfile_path = os.path.join(args.output, filename) + '.csv'
checkpoint_path = os.path.join(args.check_path, filename) + '.pth'
if not os.path.exists(outfile_path):
with open(outfile_path, mode='w') as out_file:
writer = csv.DictWriter(out_file,
fieldnames=fieldnames,
lineterminator='\n')
writer.writeheader()
epoch = 1
checkpoint = torch.load(checkpoint_path) if os.path.exists(
checkpoint_path) else None
if checkpoint is not None:
pre_model.load_state_dict(checkpoint['pre_model_state_dict'])
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
scheduler.load_state_dict(checkpoint['scheduler'])
epoch = checkpoint['epoch']
pre_model.to(device)
model.to(device)
print("*** LOADED CHECKPOINT ***"
"\n{}"
"\nSeed: {}"
"\nEpoch: {}"
"\nActfun: {}"
"\np: {}"
"\nk: {}"
"\ng: {}"
"\nperm_method: {}".format(checkpoint_path,
checkpoint['curr_seed'],
checkpoint['epoch'],
checkpoint['actfun'], checkpoint['p'],
checkpoint['k'], checkpoint['g'],
checkpoint['perm_method']))
args.mix_pre_apex = False
if args.control_amp == 'apex':
args.mix_pre_apex = True
model, optimizer = amp.initialize(model, optimizer, opt_level="O2")
# ================================================================================= Training
while epoch <= args.epochs:
if args.check_path != '':
torch.save(
{
'pre_model_state_dict': pre_model.state_dict(),
'model_state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(),
'curr_seed': args.seed,
'epoch': epoch,
'actfun': args.actfun,
'p': args.p,
'k': args.k,
'g': args.g,
'perm_method': 'shuffle'
}, checkpoint_path)
util.seed_all((args.seed * args.epochs) + epoch)
start_time = time.time()
args.mix_pre = False
if args.control_amp == 'native':
args.mix_pre = True
scaler = torch.cuda.amp.GradScaler()
# ---- Training
model.train()
total_train_loss, n, num_correct, num_total = 0, 0, 0, 0
for batch_idx, (x, targetx) in enumerate(loader_train):
x, targetx = x.to(device), targetx.to(device)
optimizer.zero_grad()
if args.mix_pre:
with torch.cuda.amp.autocast():
with torch.no_grad():
x = pre_model(x)
output = model(x)
train_loss = criterion(output, targetx)
total_train_loss += train_loss
n += 1
scaler.scale(train_loss).backward()
scaler.step(optimizer)
scaler.update()
elif args.mix_pre_apex:
with torch.no_grad():
x = pre_model(x)
output = model(x)
train_loss = criterion(output, targetx)
total_train_loss += train_loss
n += 1
with amp.scale_loss(train_loss, optimizer) as scaled_loss:
scaled_loss.backward()
optimizer.step()
else:
with torch.no_grad():
x = pre_model(x)
output = model(x)
train_loss = criterion(output, targetx)
total_train_loss += train_loss
n += 1
train_loss.backward()
optimizer.step()
scheduler.step()
_, prediction = torch.max(output.data, 1)
num_correct += torch.sum(prediction == targetx.data)
num_total += len(prediction)
epoch_aug_train_loss = total_train_loss / n
epoch_aug_train_acc = num_correct * 1.0 / num_total
alpha_primes = []
alphas = []
if model.actfun == 'combinact':
for i, layer_alpha_primes in enumerate(model.all_alpha_primes):
curr_alpha_primes = torch.mean(layer_alpha_primes, dim=0)
curr_alphas = F.softmax(curr_alpha_primes, dim=0).data.tolist()
curr_alpha_primes = curr_alpha_primes.tolist()
alpha_primes.append(curr_alpha_primes)
alphas.append(curr_alphas)
model.eval()
with torch.no_grad():
total_val_loss, n, num_correct, num_total = 0, 0, 0, 0
for batch_idx, (y, targety) in enumerate(loader_eval):
y, targety = y.to(device), targety.to(device)
with torch.no_grad():
y = pre_model(y)
output = model(y)
val_loss = criterion(output, targety)
total_val_loss += val_loss
n += 1
_, prediction = torch.max(output.data, 1)
num_correct += torch.sum(prediction == targety.data)
num_total += len(prediction)
epoch_val_loss = total_val_loss / n
epoch_val_acc = num_correct * 1.0 / num_total
lr_curr = 0
for param_group in optimizer.param_groups:
lr_curr = param_group['lr']
print(
" Epoch {}: LR {:1.5f} ||| aug_train_acc {:1.4f} | val_acc {:1.4f} ||| "
"aug_train_loss {:1.4f} | val_loss {:1.4f} ||| time = {:1.4f}".
format(epoch, lr_curr, epoch_aug_train_acc, epoch_val_acc,
epoch_aug_train_loss, epoch_val_loss,
(time.time() - start_time)),
flush=True)
epoch_train_loss = 0
epoch_train_acc = 0
if epoch == args.epochs:
with torch.no_grad():
total_train_loss, n, num_correct, num_total = 0, 0, 0, 0
for batch_idx, (x, targetx) in enumerate(loader_train):
x, targetx = x.to(device), targetx.to(device)
with torch.no_grad():
x = pre_model(x)
output = model(x)
train_loss = criterion(output, targetx)
total_train_loss += train_loss
n += 1
_, prediction = torch.max(output.data, 1)
num_correct += torch.sum(prediction == targetx.data)
num_total += len(prediction)
epoch_aug_train_loss = total_train_loss / n
epoch_aug_train_acc = num_correct * 1.0 / num_total
total_train_loss, n, num_correct, num_total = 0, 0, 0, 0
for batch_idx, (x, targetx) in enumerate(loader_eval):
x, targetx = x.to(device), targetx.to(device)
with torch.no_grad():
x = pre_model(x)
output = model(x)
train_loss = criterion(output, targetx)
total_train_loss += train_loss
n += 1
_, prediction = torch.max(output.data, 1)
num_correct += torch.sum(prediction == targetx.data)
num_total += len(prediction)
epoch_train_loss = total_val_loss / n
epoch_train_acc = num_correct * 1.0 / num_total
# Outputting data to CSV at end of epoch
with open(outfile_path, mode='a') as out_file:
writer = csv.DictWriter(out_file,
fieldnames=fieldnames,
lineterminator='\n')
writer.writerow({
'dataset': args.data,
'seed': args.seed,
'epoch': epoch,
'time': (time.time() - start_time),
'actfun': model.actfun,
'model': args.model,
'batch_size': args.batch_size,
'alpha_primes': alpha_primes,
'alphas': alphas,
'num_params': util.get_model_params(model),
'k': args.k,
'p': args.p,
'g': args.g,
'perm_method': 'shuffle',
'gen_gap': float(epoch_val_loss - epoch_train_loss),
'epoch_train_loss': float(epoch_train_loss),
'epoch_train_acc': float(epoch_train_acc),
'epoch_aug_train_loss': float(epoch_aug_train_loss),
'epoch_aug_train_acc': float(epoch_aug_train_acc),
'epoch_val_loss': float(epoch_val_loss),
'epoch_val_acc': float(epoch_val_acc),
'curr_lr': lr_curr,
'found_lr': args.lr,
'epochs': args.epochs
})
epoch += 1