def __init__(self, model, recipe):
self.enabled = bool(recipe)
self.model = model.module if is_parallel(model) else model
self.recipe = recipe
self.manager = ScheduledModifierManager.from_yaml(
recipe) if self.enabled else None
self.logger = None
def _setup_manager(
self, kwargs
) -> Tuple[Optional[ScheduledModifierManager],
List[ScheduledModifierManager]]:
manager = None
arch_managers = []
if self.recipe is not None:
manager = ScheduledModifierManager.from_yaml(
self.recipe, recipe_variables=self.recipe_args)
_LOGGER.info(
"Loaded SparseML recipe variable into manager for recipe: "
f"{self.recipe} and recipe_variables: {self.recipe_args}")
arch_recipe_paths = glob.glob(
os.path.join(self.model_state_path, RECIPE_REGEX))
if arch_recipe_paths:
arch_managers = [
ScheduledModifierManager.from_yaml(path)
for path in arch_recipe_paths
]
_LOGGER.info(
f"Loaded SparseML {len(arch_recipe_paths)} recipes into architecture "
f"managers from {arch_recipe_paths}")
if manager is not None and manager in arch_managers:
# new recipe and the one stored with model are the same,
# keep manager and remove from arch_managers to keep from applying twice.
# remove this logic once recipe stages land
arch_managers.remove(manager)
_LOGGER.info(
"Removed duplicate SparseML recipe from arch_managers that matched "
"the recipe variable to prevent double application")
if (manager is not None and manager.max_epochs and "args" in kwargs
and (hasattr(kwargs["args"], "num_train_epochs"))):
_LOGGER.warning(
f"Overriding num_train_epochs from Recipe to {manager.max_epochs}"
)
kwargs["args"].num_train_epochs = manager.max_epochs
return manager, arch_managers
def __init__(self, recipe_path):
if not _SPARSEML_AVAILABLE:
if not _PL_GREATER_EQUAL_1_4_5:
raise MisconfigurationException(
"SparseML requires PyTorch Lightning 1.4.5 or greater.")
if not _TORCH_MAX_VERSION_1_8_1:
raise MisconfigurationException(
"SparseML requires PyTorch version 1.8.1 or lower.")
raise MisconfigurationException(
"SparseML has not be installed, install with pip install sparseml"
)
self.manager = ScheduledModifierManager.from_yaml(recipe_path)
def create_scheduled_optimizer(
train_args: Any,
model: Module,
train_loader: DataLoader,
loggers: List[Any],
) -> Tuple[int, ScheduledOptimizer, ScheduledModifierManager]:
"""
:param train_args : An object with task specific config
:param model: model architecture to train
:param train_loader: A DataLoader for training data
:param loggers: List of loggers to use during training process
:type train_args: TrainingArguments
"""
# # optimizer setup
optim_const = torch.optim.__dict__[train_args.optim]
optim = optim_const(model.parameters(),
lr=train_args.init_lr,
**train_args.optim_args)
print(f"created optimizer: {optim}")
print(
"note, the lr for the optimizer may not reflect the manager yet until "
"the recipe config is created and run")
# restore from previous check point
if train_args.checkpoint_path:
# currently optimizer restoring is unsupported
# mapping of the restored params to the correct device is not working
# load_optimizer(args.checkpoint_path, optim)
epoch = 0 # load_epoch(args.checkpoint_path) + 1
print(f"restored checkpoint from {train_args.checkpoint_path} for "
f"epoch {epoch - 1}")
else:
epoch = 0
# modifier setup
add_mods = (ConstantPruningModifier.from_sparse_model(model)
if train_args.sparse_transfer_learn else None)
manager = ScheduledModifierManager.from_yaml(
file_path=train_args.recipe_path, add_modifiers=add_mods)
optim = ScheduledOptimizer(
optim,
model,
manager,
steps_per_epoch=len(train_loader),
loggers=loggers,
)
print(f"created manager: {manager}")
return epoch, optim, manager
def main(args):
############################
# logging and saving setup #
############################
save_dir = os.path.abspath(os.path.expanduser(args.save_dir))
# get unique model tag, defaults to '{model_name}'
if not args.model_tag:
model_tag = args.model.replace("/", ".")
model_id = model_tag
model_inc = 0
while os.path.exists(os.path.join(args.save_dir, model_id)):
model_inc += 1
model_id = "{}__{:02d}".format(model_tag, model_inc)
else:
model_id = args.model_tag
save_dir = os.path.join(save_dir, model_id)
create_dirs(save_dir)
print("Model id is set to {}".format(model_id))
###########################
# standard training setup #
###########################
# create data loaders
train_loader, _, _ = _create_imagefolder_dataloader(args, train=True)
val_loader, num_classes, image_shape = _create_imagefolder_dataloader(
args, train=False
)
dataloaders = {"train": train_loader, "val": val_loader}
# create model
model = _get_torchvision_model(
args.model,
num_classes,
args.pretrained,
args.checkpoint_path,
)
print("created model: {}".format(model))
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
print("using device: {}".format(device))
# create standard SGD optimizer and cross entropy loss function
criterion = CrossEntropyLoss()
optimizer = SGD(
model.parameters(), lr=0.001, momentum=0.9
) # lr will be overridden by recipe
##########################
# add sparseml modifiers #
##########################
manager = ScheduledModifierManager.from_yaml(args.recipe_path)
optimizer = ScheduledOptimizer(
optimizer,
model,
manager,
steps_per_epoch=len(train_loader),
loggers=[PythonLogger()],
)
########################
# torchvision training #
########################
model, val_acc_history = train_model(
model,
dataloaders,
criterion,
optimizer,
device,
num_epochs=manager.max_epochs,
is_inception="inception" in args.model,
)
########################
# export trained model #
########################
exporter = ModuleExporter(model, save_dir)
sample_input = torch.randn(image_shape).unsqueeze(0) # sample batch for ONNX export
exporter.export_onnx(sample_input)
exporter.export_pytorch()
print("Model ONNX export and PyTorch weights saved to {}".format(save_dir))
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))
# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# flake8: noqa
import math
from sparseml.pytorch.optim import ScheduledModifierManager, ScheduledOptimizer
from sparseml.pytorch.utils import PythonLogger, TensorBoardLogger
manager = ScheduledModifierManager.from_yaml("/PATH/TO/config.yaml")
optimizer = ScheduledOptimizer(
optimizer,
MODEL,
manager,
steps_per_epoch=math.ceil(len(TRAIN_DATASET) / TRAIN_BATCH_SIZE),
loggers=[TensorBoardLogger(), PythonLogger()],
)
def train(hyp, opt, device, tb_writer=None, wandb=None):
logger.info(
colorstr('hyperparameters: ') + ', '.join(f'{k}={v}'
for k, v in hyp.items()))
save_dir, epochs, batch_size, total_batch_size, weights, rank = \
Path(opt.save_dir), opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank
# Directories
wdir = save_dir / 'weights'
wdir.mkdir(parents=True, exist_ok=True) # make dir
last = wdir / 'last.pt'
best = wdir / 'best.pt'
results_file = save_dir / 'results.txt'
# Save run settings
with open(save_dir / 'hyp.yaml', 'w') as f:
yaml.dump(hyp, f, sort_keys=False)
with open(save_dir / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False)
# Configure
plots = not opt.evolve # create plots
cuda = device.type != 'cpu'
init_seeds(2 + rank)
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict
with torch_distributed_zero_first(rank):
check_dataset(data_dict) # check
train_path = data_dict['train']
test_path = data_dict['val']
nc = 1 if opt.single_cls else int(data_dict['nc']) # number of classes
names = ['item'] if opt.single_cls and len(
data_dict['names']) != 1 else data_dict['names'] # class names
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (
len(names), nc, opt.data) # check
# Model
pretrained = weights.endswith('.pt') or weights.endswith(
'.pth') # SparseML integration
if pretrained:
with torch_distributed_zero_first(rank):
attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location=device) # load checkpoint
if hyp.get('anchors'):
ckpt['model'].yaml['anchors'] = round(
hyp['anchors']) # force autoanchor
model = Model(opt.cfg or ckpt['model'].yaml, ch=3,
nc=nc).to(device) # create
exclude = ['anchor'] if opt.cfg or hyp.get('anchors') else [
] # exclude keys
state_dict = _load_checkpoint_model_state_dict(
ckpt) # SparseML integration
state_dict = intersect_dicts(state_dict,
model.state_dict(),
exclude=exclude) # intersect
model.load_state_dict(state_dict, strict=False) # load
logger.info(
'Transferred %g/%g items from %s' %
(len(state_dict), len(model.state_dict()), weights)) # report
else:
model = Model(opt.cfg, ch=3, nc=nc).to(device) # create
# Freeze
freeze = [] # parameter names to freeze (full or partial)
for k, v in model.named_parameters():
v.requires_grad = True # train all layers
if any(x in k for x in freeze):
print('freezing %s' % k)
v.requires_grad = False
# Optimizer
nbs = 64 # nominal batch size
accumulate = max(round(nbs / total_batch_size),
1) # accumulate loss before optimizing
hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay
logger.info(f"Scaled weight_decay = {hyp['weight_decay']}")
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_modules():
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
pg2.append(v.bias) # biases
if isinstance(v, nn.BatchNorm2d):
pg0.append(v.weight) # no decay
elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
pg1.append(v.weight) # apply decay
if opt.adam:
optimizer = optim.Adam(pg0,
lr=hyp['lr0'],
betas=(hyp['momentum'],
0.999)) # adjust beta1 to momentum
else:
optimizer = optim.SGD(pg0,
lr=hyp['lr0'],
momentum=hyp['momentum'],
nesterov=True)
optimizer.add_param_group({
'params': pg1,
'weight_decay': hyp['weight_decay']
}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' %
(len(pg2), len(pg1), len(pg0)))
del pg0, pg1, pg2
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
if opt.linear_lr:
lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp[
'lrf'] # linear
else:
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# plot_lr_scheduler(optimizer, scheduler, epochs)
# Logging
if rank in [-1, 0] and wandb and wandb.run is None:
opt.hyp = hyp # add hyperparameters
wandb_run = wandb.init(
config=opt,
resume="allow",
project='YOLOv5'
if opt.project == 'runs/train' else Path(opt.project).stem,
name=save_dir.stem,
id=ckpt.get('wandb_id') if 'ckpt' in locals() else None)
loggers = {'wandb': wandb} # loggers dict
# Resume
start_epoch, best_fitness = 0, 0.0
if pretrained:
# Optimizer
if ckpt['optimizer'] is not None:
optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness']
# Results
if ckpt.get('training_results') is not None:
with open(results_file, 'w') as file:
file.write(ckpt['training_results']) # write results.txt
# Epochs
start_epoch = ckpt['epoch'] + 1
if opt.resume:
assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (
weights, epochs)
if epochs < start_epoch:
logger.info(
'%s has been trained for %g epochs. Fine-tuning for %g additional epochs.'
% (weights, ckpt['epoch'], epochs))
epochs += ckpt['epoch'] # finetune additional epochs
del ckpt, state_dict
# Image sizes
gs = int(model.stride.max()) # grid size (max stride)
nl = model.model[
-1].nl # number of detection layers (used for scaling hyp['obj'])
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size
] # verify imgsz are gs-multiples
# DP mode
if cuda and rank == -1 and torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
# SyncBatchNorm
if opt.sync_bn and cuda and rank != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
logger.info('Using SyncBatchNorm()')
# EMA
####################################################################################
# Start SparseML Integration - optional EMA
####################################################################################
ema = ModelEMA(model) if rank in [-1, 0] and opt.use_ema else None
####################################################################################
# End SparseML Integration - optional EMA
####################################################################################
# DDP mode
if cuda and rank != -1:
model = DDP(model,
device_ids=[opt.local_rank],
output_device=opt.local_rank)
# Trainloader
dataloader, dataset = create_dataloader(train_path,
imgsz,
batch_size,
gs,
opt,
hyp=hyp,
augment=True,
cache=opt.cache_images,
rect=opt.rect,
rank=rank,
world_size=opt.world_size,
workers=opt.workers,
image_weights=opt.image_weights,
quad=opt.quad,
prefix=colorstr('train: '))
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(dataloader) # number of batches
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (
mlc, nc, opt.data, nc - 1)
# Process 0
if rank in [-1, 0]:
if ema:
ema.updates = start_epoch * nb // accumulate # set EMA updates
testloader = create_dataloader(
test_path,
imgsz_test,
batch_size * 2,
gs,
opt, # testloader
hyp=hyp,
cache=opt.cache_images and not opt.notest,
rect=True,
rank=-1,
world_size=opt.world_size,
workers=opt.workers,
pad=0.5,
prefix=colorstr('val: '))[0]
if not opt.resume:
labels = np.concatenate(dataset.labels, 0)
c = torch.tensor(labels[:, 0]) # classes
# cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency
# model._initialize_biases(cf.to(device))
if plots:
plot_labels(labels, save_dir, loggers)
if tb_writer:
tb_writer.add_histogram('classes', c, 0)
# Anchors
if not opt.noautoanchor:
check_anchors(dataset,
model=model,
thr=hyp['anchor_t'],
imgsz=imgsz)
# Model parameters
hyp['box'] *= 3. / nl # scale to layers
hyp['cls'] *= nc / 80. * 3. / nl # scale to classes and layers
hyp['obj'] *= (imgsz / 640)**2 * 3. / nl # scale to image size and layers
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou)
model.class_weights = labels_to_class_weights(
dataset.labels, nc).to(device) * nc # attach class weights
model.names = names
####################################################################################
# Start SparseML Integration
####################################################################################
from sparseml.pytorch.nn import replace_activations
from sparseml.pytorch.optim import ScheduledModifierManager, ScheduledOptimizer
from sparseml.pytorch.utils import is_parallel_model, PythonLogger, TensorBoardLogger
if not opt.no_leaky_relu_override: # use LeakyReLU activations
model = replace_activations(model, "lrelu", inplace=True)
manager = ScheduledModifierManager.from_yaml(opt.sparseml_recipe)
optimizer = ScheduledOptimizer(
optimizer,
model if not is_parallel_model(model) else model.module,
manager,
steps_per_epoch=len(dataloader),
loggers=[PythonLogger(),
TensorBoardLogger(writer=tb_writer)])
# override lr scheduler if recipe makes any LR updates
if any("LearningRate" in str(modifier) for modifier in manager.modifiers):
logger.info(
"Disabling yolo LR scheduler, managing LR using SparseML recipe")
scheduler = None
# disable model pickling if QAT is set
qat = False
if any("Quantization" in str(modifier) for modifier in manager.modifiers):
logger.info("Disabling pickling for Yolo model, QAT modifiers present")
qat = True
if manager.max_epochs:
epochs = manager.max_epochs or epochs # override num_epochs
logger.info(
f"overriding number of epochs from SparseML manager to {manager.max_epochs}"
)
####################################################################################
# End SparseML Integration
####################################################################################
# Start training
t0 = time.time()
nw = max(round(hyp['warmup_epochs'] * nb),
1000) # number of warmup iterations, max(3 epochs, 1k iterations)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
maps = np.zeros(nc) # mAP per class
results = (0, 0, 0, 0, 0, 0, 0
) # P, R, [email protected], [email protected], val_loss(box, obj, cls)
if scheduler: # SparseML integration
scheduler.last_epoch = start_epoch - 1 # do not move
scaler = amp.GradScaler(enabled=(cuda and opt.use_amp))
compute_loss = ComputeLoss(model) # init loss class
logger.info(f'Image sizes {imgsz} train, {imgsz_test} test\n'
f'Using {dataloader.num_workers} dataloader workers\n'
f'Logging results to {save_dir}\n'
f'Starting training for {epochs} epochs...')
for epoch in range(
start_epoch, epochs
): # epoch ------------------------------------------------------------------
model.train()
# Update image weights (optional)
if opt.image_weights:
# Generate indices
if rank in [-1, 0]:
cw = model.class_weights.cpu().numpy() * (
1 - maps)**2 / nc # class weights
iw = labels_to_image_weights(dataset.labels,
nc=nc,
class_weights=cw) # image weights
dataset.indices = random.choices(
range(dataset.n), weights=iw,
k=dataset.n) # rand weighted idx
# Broadcast if DDP
if rank != -1:
indices = (torch.tensor(dataset.indices)
if rank == 0 else torch.zeros(dataset.n)).int()
dist.broadcast(indices, 0)
if rank != 0:
dataset.indices = indices.cpu().numpy()
# Update mosaic border
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(4, device=device) # mean losses
if rank != -1:
dataloader.sampler.set_epoch(epoch)
pbar = enumerate(dataloader)
logger.info(
('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls',
'total', 'targets', 'img_size'))
if rank in [-1, 0]:
pbar = tqdm(pbar, total=nb) # progress bar
optimizer.zero_grad()
for i, (
imgs, targets, paths, _
) in pbar: # batch -------------------------------------------------------------
ni = i + nb * epoch # number integrated batches (since train start)
imgs = imgs.to(device, non_blocking=True).float(
) / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
# Warmup
if ni <= nw:
xi = [0, nw] # x interp
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
accumulate = max(
1,
np.interp(ni, xi, [1, nbs / total_batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
if scheduler: # SparseML integration, do not force warmup lr when overriding
x['lr'] = np.interp(ni, xi, [
hyp['warmup_bias_lr'] if j == 2 else 0.0,
x['initial_lr'] * lf(epoch)
])
if 'momentum' in x:
x['momentum'] = np.interp(
ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5,
imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]
] # new shape (stretched to gs-multiple)
imgs = F.interpolate(imgs,
size=ns,
mode='bilinear',
align_corners=False)
# Forward
with amp.autocast(enabled=(cuda and opt.use_amp)):
pred = model(imgs) # forward
loss, loss_items = compute_loss(
pred, targets.to(device)) # loss scaled by batch_size
if rank != -1:
loss *= opt.world_size # gradient averaged between devices in DDP mode
if opt.quad:
loss *= 4.
# Backward
scaler.scale(loss).backward()
# Optimize
if ni % accumulate == 0:
scaler.step(optimizer) # optimizer.step
scaler.update()
optimizer.zero_grad()
if ema:
ema.update(model)
# Print
if rank in [-1, 0]:
mloss = (mloss * i + loss_items) / (i + 1
) # update mean losses
mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9
if torch.cuda.is_available() else 0) # (GB)
s = ('%10s' * 2 +
'%10.4g' * 6) % ('%g/%g' % (epoch, epochs - 1), mem,
*mloss, targets.shape[0], imgs.shape[-1])
pbar.set_description(s)
# Plot
if plots and ni < 3:
f = save_dir / f'train_batch{ni}.jpg' # filename
Thread(target=plot_images,
args=(imgs, targets, paths, f),
daemon=True).start()
# if tb_writer:
# tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
# tb_writer.add_graph(model, imgs) # add model to tensorboard
elif plots and ni == 10 and wandb:
wandb.log(
{
"Mosaics": [
wandb.Image(str(x), caption=x.name)
for x in save_dir.glob('train*.jpg')
if x.exists()
]
},
commit=False)
# end batch ------------------------------------------------------------------------------------------------
# end epoch ----------------------------------------------------------------------------------------------------
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for tensorboard
if scheduler: # SparseML integration
scheduler.step()
# DDP process 0 or single-GPU
if rank in [-1, 0]:
# mAP
if ema:
ema.update_attr(model,
include=[
'yaml', 'nc', 'hyp', 'gr', 'names',
'stride', 'class_weights'
])
final_epoch = epoch + 1 == epochs
if not opt.notest or final_epoch: # Calculate mAP
results, maps, times = test.test(
opt.data,
batch_size=batch_size * 2,
imgsz=imgsz_test,
model=ema.ema if ema else model,
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=save_dir,
verbose=nc < 50 and final_epoch,
plots=plots and final_epoch,
log_imgs=opt.log_imgs if wandb else 0,
compute_loss=compute_loss,
half_precision=opt.use_amp) # SparseML integration
# Write
with open(results_file, 'a') as f:
f.write(
s + '%10.4g' * 7 % results +
'\n') # P, R, [email protected], [email protected], val_loss(box, obj, cls)
if len(opt.name) and opt.bucket:
os.system('gsutil cp %s gs://%s/results/results%s.txt' %
(results_file, opt.bucket, opt.name))
# Log
tags = [
'train/box_loss',
'train/obj_loss',
'train/cls_loss', # train loss
'metrics/precision',
'metrics/recall',
'metrics/mAP_0.5',
'metrics/mAP_0.5:0.95',
'val/box_loss',
'val/obj_loss',
'val/cls_loss', # val loss
'x/lr0',
'x/lr1',
'x/lr2'
] # params
for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags):
if tb_writer:
tb_writer.add_scalar(tag, x, epoch) # tensorboard
if wandb:
wandb.log({tag: x}, step=epoch,
commit=tag == tags[-1]) # W&B
# Update best mAP
fi = fitness(np.array(results).reshape(
1, -1)) # weighted combination of [P, R, [email protected], [email protected]]
if fi > best_fitness:
best_fitness = fi
# Save model
save = (not opt.nosave) or (final_epoch and not opt.evolve)
if save:
with open(results_file, 'r') as f: # create checkpoint
ckpt_model = ema.ema if ema else model if not qat else model.state_dict(
) # SparseML integration
ckpt = {
'epoch':
epoch,
'best_fitness':
best_fitness,
'training_results':
f.read(),
'model':
ckpt_model, # SparseML integration
'optimizer':
None if final_epoch else optimizer.state_dict(),
'wandb_id':
wandb_run.id if wandb else None
}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
del ckpt
# end epoch ----------------------------------------------------------------------------------------------------
# end training
if rank in [-1, 0]:
# Strip optimizers
final = best if best.exists() else last # final model
for f in [last, best]:
if f.exists(
) and not qat: # SparseML integration - qat state dict incompatible
strip_optimizer(f) # strip optimizers
if opt.bucket:
os.system(f'gsutil cp {final} gs://{opt.bucket}/weights') # upload
# Plots
if plots:
plot_results(save_dir=save_dir) # save as results.png
if wandb:
files = [
'results.png', 'confusion_matrix.png',
*[f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R')]
]
wandb.log({
"Results": [
wandb.Image(str(save_dir / f), caption=f)
for f in files if (save_dir / f).exists()
]
})
if opt.log_artifacts:
wandb.log_artifact(artifact_or_path=str(final),
type='model',
name=save_dir.stem)
# Test best.pt
logger.info('%g epochs completed in %.3f hours.\n' %
(epoch - start_epoch + 1, (time.time() - t0) / 3600))
if opt.data.endswith('coco.yaml') and nc == 80: # if COCO
for conf, iou, save_json in ([0.25, 0.45,
False], [0.001, 0.65,
True]): # speed, mAP tests
# SparseML integration - load test model
test_model = model if qat else attempt_load(final, device)
if opt.use_amp:
test_model = test_model.half()
results, _, _ = test.test(
opt.data,
batch_size=batch_size * 2,
imgsz=imgsz_test,
conf_thres=conf,
iou_thres=iou,
model=test_model,
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=save_dir,
save_json=save_json,
plots=False,
half_precision=opt.use_amp) # SparseML integration
#################################################################################
# Start SparseML ONNX Export
#################################################################################
from sparseml.pytorch.utils import ModuleExporter
from sparseml.pytorch.utils.quantization import skip_onnx_input_quantize
onnx_path = f"{save_dir}/model.onnx"
logger.info(f"training complete, exporting ONNX to {onnx_path}")
export_model = model.module if is_parallel_model(model) else model
export_model.model[
-1].export = True # do not export grid post-procesing
exporter = ModuleExporter(export_model, save_dir)
exporter.export_onnx(torch.randn(1, 3, imgsz, imgsz),
convert_qat=True)
if qat:
skip_onnx_input_quantize(onnx_path, onnx_path)
#################################################################################
# End SparseML ONNX Export
#################################################################################
else:
dist.destroy_process_group()
wandb.run.finish() if wandb and wandb.run else None
torch.cuda.empty_cache()
return results
def test_oracle_recipe_application(oracle_test_params):
manager = ScheduledModifierManager.from_yaml(oracle_test_params.generated_recipe)
model = oracle_test_params.model_lambda()
manager.apply(model)
assert model is not None
def train(
working_dir: str,
config_path: str,
model: Module,
train_dataset: Dataset,
val_dataset: Dataset,
batch_size: int,
optim_const: Callable[[Module], Optimizer],
loss: Union[LossWrapper, Callable[[Any, Any], Tensor]],
devices: str,
):
"""
Dataset setup
"""
LOGGER.info("batch_size set to {}".format(batch_size))
LOGGER.info("train_dataset set to {}".format(train_dataset))
LOGGER.info("val_dataset set to {}".format(val_dataset))
train_loader = DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=8,
pin_memory=True,
)
val_loader = DataLoader(
val_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=8,
pin_memory=True,
)
"""
Model, optimizer, loss setup
"""
model_dir = clean_path(os.path.join(working_dir, "model"))
optim = optim_const(model)
LOGGER.info("model set to {}".format(model))
LOGGER.info("optimizer set to {}".format(optim))
LOGGER.info("loss set to {}".format(loss))
LOGGER.info("devices set to {}".format(devices))
"""
Manager and config setup
"""
manager = ScheduledModifierManager.from_yaml(config_path)
logs_dir = clean_path(os.path.join(working_dir, "logs"))
loggers = [TensorBoardLogger(logs_dir), PythonLogger()]
optim = ScheduledOptimizer(
optim, model, manager, steps_per_epoch=len(train_loader), loggers=loggers
)
"""
Training and testing
"""
model, device, device_ids = model_to_device(model, devices)
trainer = ModuleTrainer(model, device, loss, optim, loggers=loggers)
tester = ModuleTester(model, device, loss, loggers=loggers, log_steps=-1)
epoch = -1
tester.run_epoch(val_loader, epoch=epoch)
for epoch in range(manager.max_epochs):
LOGGER.info("starting training epoch {}".format(epoch))
train_res = trainer.run_epoch(train_loader, epoch)
LOGGER.info("finished training epoch {}: {}".format(epoch, train_res))
val_res = tester.run_epoch(val_loader, epoch)
LOGGER.info("finished validation epoch {}: {}".format(epoch, val_res))
exporter = ModuleExporter(model, model_dir)
exporter.export_pytorch(optim, epoch)
for data in val_loader:
exporter.export_onnx(data)
