scripts/train_detection.py [182:283]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - def make_optimizer(config, model, epoch_length): body_config = config["body"] opt_config = config["optimizer"] sch_config = config["scheduler"] # Gather parameters from the network norm_parameters = [] other_parameters = [] for m in model.modules(): if any(isinstance(m, layer) for layer in NORM_LAYERS): norm_parameters += [p for p in m.parameters() if p.requires_grad] elif any(isinstance(m, layer) for layer in OTHER_LAYERS): other_parameters += [p for p in m.parameters() if p.requires_grad] assert len(norm_parameters) + len(other_parameters) == len([p for p in model.parameters() if p.requires_grad]), \ "Not all parameters that require grad are accounted for in the optimizer" # Set-up optimizer hyper-parameters parameters = [ { "params": norm_parameters, "lr": opt_config.getfloat("lr") if not body_config.getboolean("bn_frozen") else 0., "weight_decay": opt_config.getfloat("weight_decay") if opt_config.getboolean("weight_decay_norm") else 0. }, { "params": other_parameters, "lr": opt_config.getfloat("lr"), "weight_decay": opt_config.getfloat("weight_decay") } ] optimizer = optim.SGD( parameters, momentum=opt_config.getfloat("momentum"), nesterov=opt_config.getboolean("nesterov")) scheduler = scheduler_from_config(sch_config, optimizer, epoch_length) assert sch_config["update_mode"] in ("batch", "epoch") batch_update = sch_config["update_mode"] == "batch" total_epochs = sch_config.getint("epochs") return optimizer, scheduler, batch_update, total_epochs def train(model, optimizer, scheduler, dataloader, meters, **varargs): model.train() dataloader.batch_sampler.set_epoch(varargs["epoch"]) optimizer.zero_grad() global_step = varargs["global_step"] loss_weights = varargs["loss_weights"] data_time_meter = AverageMeter((), meters["loss"].momentum) batch_time_meter = AverageMeter((), meters["loss"].momentum) data_time = time.time() for it, batch in enumerate(dataloader): # Upload batch batch = {k: batch[k].cuda(device=varargs["device"], non_blocking=True) for k in NETWORK_INPUTS} data_time_meter.update(torch.tensor(time.time() - data_time)) # Update scheduler global_step += 1 if varargs["batch_update"]: scheduler.step(global_step) batch_time = time.time() # Run network losses, _ = model(**batch, do_loss=True, do_prediction=False) distributed.barrier() losses = OrderedDict((k, v.mean()) for k, v in losses.items()) losses["loss"] = sum(w * l for w, l in zip(loss_weights, losses.values())) optimizer.zero_grad() losses["loss"].backward() optimizer.step() # Gather stats from all workers losses = all_reduce_losses(losses) # Update meters with torch.no_grad(): for loss_name, loss_value in losses.items(): meters[loss_name].update(loss_value.cpu()) batch_time_meter.update(torch.tensor(time.time() - batch_time)) # Clean-up del batch, losses # Log if varargs["summary"] is not None and (it + 1) % varargs["log_interval"] == 0: logging.iteration( varargs["summary"], "train", global_step, varargs["epoch"] + 1, varargs["num_epochs"], it + 1, len(dataloader), OrderedDict([ ("lr", scheduler.get_lr()[0]), ("loss", meters["loss"]), ("obj_loss", meters["obj_loss"]), ("bbx_loss", meters["bbx_loss"]), ("roi_cls_loss", meters["roi_cls_loss"]), ("roi_bbx_loss", meters["roi_bbx_loss"]), - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - scripts/train_instance_seg.py [188:289]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - def make_optimizer(config, model, epoch_length): body_config = config["body"] opt_config = config["optimizer"] sch_config = config["scheduler"] # Gather parameters from the network norm_parameters = [] other_parameters = [] for m in model.modules(): if any(isinstance(m, layer) for layer in NORM_LAYERS): norm_parameters += [p for p in m.parameters() if p.requires_grad] elif any(isinstance(m, layer) for layer in OTHER_LAYERS): other_parameters += [p for p in m.parameters() if p.requires_grad] assert len(norm_parameters) + len(other_parameters) == len([p for p in model.parameters() if p.requires_grad]), \ "Not all parameters that require grad are accounted for in the optimizer" # Set-up optimizer hyper-parameters parameters = [ { "params": norm_parameters, "lr": opt_config.getfloat("lr") if not body_config.getboolean("bn_frozen") else 0., "weight_decay": opt_config.getfloat("weight_decay") if opt_config.getboolean("weight_decay_norm") else 0. }, { "params": other_parameters, "lr": opt_config.getfloat("lr"), "weight_decay": opt_config.getfloat("weight_decay") } ] optimizer = optim.SGD( parameters, momentum=opt_config.getfloat("momentum"), nesterov=opt_config.getboolean("nesterov")) scheduler = scheduler_from_config(sch_config, optimizer, epoch_length) assert sch_config["update_mode"] in ("batch", "epoch") batch_update = sch_config["update_mode"] == "batch" total_epochs = sch_config.getint("epochs") return optimizer, scheduler, batch_update, total_epochs def train(model, optimizer, scheduler, dataloader, meters, **varargs): model.train() dataloader.batch_sampler.set_epoch(varargs["epoch"]) optimizer.zero_grad() global_step = varargs["global_step"] loss_weights = varargs["loss_weights"] data_time_meter = AverageMeter((), meters["loss"].momentum) batch_time_meter = AverageMeter((), meters["loss"].momentum) data_time = time.time() for it, batch in enumerate(dataloader): # Upload batch batch = {k: batch[k].cuda(device=varargs["device"], non_blocking=True) for k in NETWORK_INPUTS} data_time_meter.update(torch.tensor(time.time() - data_time)) # Update scheduler global_step += 1 if varargs["batch_update"]: scheduler.step(global_step) batch_time = time.time() # Run network losses, _ = model(**batch, do_loss=True, do_prediction=False) distributed.barrier() losses = OrderedDict((k, v.mean()) for k, v in losses.items()) losses["loss"] = sum(w * l for w, l in zip(loss_weights, losses.values())) optimizer.zero_grad() losses["loss"].backward() optimizer.step() # Gather stats from all workers losses = all_reduce_losses(losses) # Update meters with torch.no_grad(): for loss_name, loss_value in losses.items(): meters[loss_name].update(loss_value.cpu()) batch_time_meter.update(torch.tensor(time.time() - batch_time)) # Clean-up del batch, losses # Log if varargs["summary"] is not None and (it + 1) % varargs["log_interval"] == 0: logging.iteration( varargs["summary"], "train", global_step, varargs["epoch"] + 1, varargs["num_epochs"], it + 1, len(dataloader), OrderedDict([ ("lr", scheduler.get_lr()[0]), ("loss", meters["loss"]), ("obj_loss", meters["obj_loss"]), ("bbx_loss", meters["bbx_loss"]), ("roi_cls_loss", meters["roi_cls_loss"]), ("roi_bbx_loss", meters["roi_bbx_loss"]), - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -