def wrap(model): if world_size == 1: return model else: return DistributedDataParallel( model, # find_unused_parameters=True, device_ids=[torch.cuda.current_device()], output_device=torch.cuda.current_device())
def block_backward_sync(self, model: DistributedDataParallel): """ Blocks ddp sync gradients behaviour on backwards pass. This is useful for skipping sync when accumulating gradients, reducing communication overhead Returns: context manager with sync behaviour off """ yield model.no_sync()
def _setup_model(self, model: Module) -> DistributedDataParallel: """Wraps the model into a :class:`~torch.nn.parallel.distributed.DistributedDataParallel` module.""" device_ids = self.determine_ddp_device_ids() log.detail( f"setting up DDP model with device ids: {device_ids}, kwargs: {self._ddp_kwargs}" ) return DistributedDataParallel(module=model, device_ids=device_ids, **self._ddp_kwargs)
def __init__(self, fc_dqn): super().__init__() self.state_dim = fc_dqn.state_dim self.action_dim = fc_dqn.action_dim current_device = torch.cuda.current_device() self.data_parallel = DistributedDataParallel( fc_dqn.fc, device_ids=[current_device], output_device=current_device ) self.fc_dqn = fc_dqn
def __init__(self, seq2slate_net: Seq2SlateNet): super().__init__() current_device = torch.cuda.current_device() self.data_parallel = DistributedDataParallel( seq2slate_net.seq2slate, device_ids=[current_device], output_device=current_device, ) self.seq2slate_net = seq2slate_net
def run(rank, model, train_dataloader, num_epochs, precision, accelerator, tmpdir): os.environ["LOCAL_RANK"] = str(rank) if torch.distributed.is_available() and not torch.distributed.is_initialized(): torch.distributed.init_process_group("gloo", rank=rank, world_size=2) to_device = partial(move_data_to_device, device=torch.device("cuda", rank)) model = DistributedDataParallel( to_device(model), device_ids=[rank], ) train_dataloader = DataLoader( train_dataloader.dataset, sampler=DistributedSampler(train_dataloader.dataset, rank=rank, num_replicas=2, seed=42, drop_last=False), ) with precision_context(precision, accelerator): main(to_device, model, train_dataloader, num_epochs=num_epochs) if rank == 0: atomic_save(model.state_dict(), os.path.join(tmpdir, "model_spawn.pt"))
def load_vaes_256(H, logprint): vae = VAE_256(H) if H.restore_path: logprint(f'Restoring vae from {H.restore_path}') restore_params(vae, H.restore_path, map_cpu=True, local_rank=H.local_rank, mpi_size=H.mpi_size) ema_vae = VAE_256(H) if H.restore_ema_path: logprint(f'Restoring ema vae from {H.restore_ema_path}') restore_params(ema_vae, H.restore_ema_path, map_cpu=True, local_rank=H.local_rank, mpi_size=H.mpi_size) else: ema_vae.load_state_dict(vae.state_dict()) ema_vae.requires_grad_(False) vae = vae.cuda(H.local_rank) ema_vae = ema_vae.cuda(H.local_rank) if H.image_size == 64: vae.decoder.requires_grad_(False) if H.image_size == 256: vae.encoder.requires_grad_(False) vae = DistributedDataParallel(vae, device_ids=[H.local_rank], output_device=H.local_rank, find_unused_parameters=True) if len(list(vae.named_parameters())) != len(list(vae.parameters())): raise ValueError('Some params are not named. Please name all params.') total_params = 0 for name, p in vae.named_parameters(): total_params += np.prod(p.shape) logprint(total_params=total_params, readable=f'{total_params:,}') return vae, ema_vae
def main(rank: int, epochs: int, model: nn.Module, train_loader: DataLoader, test_loader: DataLoader) -> nn.Module: device = torch.device(f'cuda:{rank}') model = model.to(device) model = DistributedDataParallel(model, device_ids=[rank], output_device=rank) # initialize optimizer and loss function optimizer = optim.SGD(model.parameters(), lr=0.01) loss = nn.CrossEntropyLoss() # train the model for i in range(epochs): model.train() train_loader.sampler.set_epoch(i) epoch_loss = 0 # train the model for one epoch pbar = tqdm(train_loader) for x, y in pbar: x = x.to(device, non_blocking=True) y = y.to(device, non_blocking=True) x = x.view(x.shape[0], -1) optimizer.zero_grad() y_hat = model(x) batch_loss = loss(y_hat, y) batch_loss.backward() optimizer.step() batch_loss_scalar = batch_loss.item() epoch_loss += batch_loss_scalar / x.shape[0] pbar.set_description(f'training batch_loss={batch_loss_scalar:.4f}') # calculate validation loss with torch.no_grad(): model.eval() val_loss = 0 pbar = tqdm(test_loader) for x, y in pbar: x = x.to(device, non_blocking=True) y = y.to(device, non_blocking=True) x = x.view(x.shape[0], -1) y_hat = model(x) batch_loss = loss(y_hat, y) batch_loss_scalar = batch_loss.item() val_loss += batch_loss_scalar / x.shape[0] pbar.set_description(f'validation batch_loss={batch_loss_scalar:.4f}') print(f"Epoch={i}, train_loss={epoch_loss:.4f}, val_loss={val_loss:.4f}") return model.module
def wrap(model): """wrap. Wrap a model with DDP if distributed training is enabled. """ if world_size == 1: return model else: return DistributedDataParallel( model, device_ids=[torch.cuda.current_device()], output_device=torch.cuda.current_device())
def load_model(self): # Build model if distutils.is_master(): print("### Loading model: {}".format(self.config["model"])) # TODO(abhshkdz): Eventually move towards computing features on-the-fly # and remove dependence from `.edge_attr`. bond_feat_dim = None if self.config["task"]["dataset"] in [ "trajectory_lmdb", "single_point_lmdb", ]: bond_feat_dim = self.config["model_attributes"].get( "num_gaussians", 50) else: raise NotImplementedError self.model = registry.get_model_class(self.config["model"])( self.train_loader.dataset[0].x.shape[-1] if hasattr(self.train_loader.dataset[0], "x") and self.train_loader.dataset[0].x is not None else None, bond_feat_dim, self.num_targets, **self.config["model_attributes"], ).to(self.device) if distutils.is_master(): print("### Loaded {} with {} parameters.".format( self.model.__class__.__name__, self.model.num_params)) if self.logger is not None: self.logger.watch(self.model) self.model = OCPDataParallel( self.model, output_device=self.device, num_gpus=1, ) if distutils.initialized(): self.model = DistributedDataParallel(self.model, device_ids=[self.device])
def validate(ae_without_ddp, data_loader, config, device, distributed, device_ids): input_shape = config['input_shape'] extended_model, _ = ae_util.get_extended_model(ae_without_ddp, config, input_shape, device, True) if distributed: extended_model = DistributedDataParallel(extended_model, device_ids=device_ids) return evaluate(extended_model, data_loader, device, split_name='Validation')
def __init__(self, config: BaseConfig): self._config = config self._model = DeepLab(num_classes=9, output_stride=8, sync_bn=False).to(self._config.device) self._border_loss = TotalLoss(self._config) self._direction_loss = CrossEntropyLoss() self._loaders = get_data_loaders(config) self._writer = SummaryWriter() self._optimizer = torch.optim.SGD(self._model.parameters(), lr=self._config.lr, weight_decay=1e-4, nesterov=True, momentum=0.9) self._scheduler = torch.optim.lr_scheduler.ExponentialLR( self._optimizer, gamma=0.97) if self._config.parallel: self._model = DistributedDataParallel(self._model, device_ids=[ self._config.device, ])
def __init__(self, cem_planner: CEMPlanner): super().__init__() self.plan_horizon_length = cem_planner.plan_horizon_length self.state_dim = cem_planner.state_dim self.action_dim = cem_planner.action_dim self.discrete_action = cem_planner.discrete_action current_device = torch.cuda.current_device() # type: ignore self.data_parallel = DistributedDataParallel( cem_planner.cem_planner_network, device_ids=[current_device], output_device=current_device, ) self.cem_planner = cem_planner
def __init__(self, mem_net): super().__init__() self.num_hiddens = mem_net.num_hiddens self.num_hidden_layers = mem_net.num_hidden_layers self.state_dim = mem_net.state_dim self.action_dim = mem_net.action_dim self.num_gaussians = mem_net.num_gaussians current_device = torch.cuda.current_device() self.data_parallel = DistributedDataParallel( mem_net.mdnrnn, device_ids=[current_device], output_device=current_device) self.mem_net = mem_net
def main(args): config = yaml_util.load_yaml_file(args.config) if args.json is not None: main_util.overwrite_config(config, args.json) distributed, device_ids = main_util.init_distributed_mode( args.world_size, args.dist_url) device = torch.device(args.device if torch.cuda.is_available() else 'cpu') teacher_model = get_model(config['teacher_model'], device) module_util.freeze_module_params(teacher_model) student_model_config = config['student_model'] student_model = get_model(student_model_config, device, require_weights=args.require_weights) reset_unfrozen = False if 'reset_unfrozen' in student_model_config: reset_unfrozen = student_model_config['reset_unfrozen'] freeze_modules(student_model, student_model_config, reset_unfrozen=reset_unfrozen) set_bottleneck_transformer(student_model, True) post_bn = False if 'post_batch_norm' in config['train']: post_bn = config['train']['post_batch_norm'] # print('Updatable parameters: {}'.format(module_util.get_updatable_param_names(student_model))) distill_backbone_only = student_model_config['distill_backbone_only'] train_config = config['train'] train_sampler, train_data_loader, val_data_loader, test_data_loader = \ data_util.get_coco_data_loaders(config['dataset'], train_config['batch_size'], distributed) if distributed: teacher_model = DataParallel(teacher_model, device_ids=device_ids) student_model = DistributedDataParallel(student_model, device_ids=device_ids) if args.distill: distill(teacher_model, student_model, train_sampler, train_data_loader, val_data_loader, device, distributed, distill_backbone_only, config, args) load_ckpt( config['student_model']['ckpt'], model=student_model.module if isinstance( student_model, DistributedDataParallel) else student_model) evaluate(teacher_model, student_model, test_data_loader, train_data_loader, device, args.skip_teacher_eval, args.transform_bottleneck, student_model_config, post_bn)
def train(model, train_loader, valid_loader, best_valid_acc, criterion, device, distributed, device_ids, train_config, num_epochs, start_epoch, init_lr, ckpt_file_path, model_type): model_without_ddp = model if distributed: model = DistributedDataParallel(model_without_ddp, device_ids=device_ids) elif device.type == 'cuda': model = DataParallel(model_without_ddp) optim_config = train_config['optimizer'] if init_lr is not None: optim_config['params']['lr'] = init_lr optimizer = func_util.get_optimizer(model, optim_config['type'], optim_config['params']) scheduler_config = train_config['scheduler'] scheduler = func_util.get_scheduler(optimizer, scheduler_config['type'], scheduler_config['params']) interval = train_config['interval'] if interval <= 0: num_batches = len(train_loader) interval = num_batches // 20 if num_batches >= 20 else 1 end_epoch = start_epoch + train_config[ 'epoch'] if num_epochs is None else start_epoch + num_epochs start_time = time.time() for epoch in range(start_epoch, end_epoch): if distributed: train_loader.sampler.set_epoch(epoch) train_epoch(model, train_loader, optimizer, criterion, epoch, device, interval) valid_acc = validate(model, valid_loader, device) if valid_acc > best_valid_acc and main_util.is_main_process(): print( 'Updating ckpt (Best top1 accuracy: {:.4f} -> {:.4f})'.format( best_valid_acc, valid_acc)) best_valid_acc = valid_acc save_ckpt(model_without_ddp, best_valid_acc, epoch, ckpt_file_path, model_type) scheduler.step() dist.barrier() total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str))
def train(epochs: int, train_data_loader: DataLoader, valid_data_loader: DataLoader = None, rank=None): device = torch.device(f'cuda:{rank}') model = create_model(model_type).to(device) model = DistributedDataParallel(model, device_ids=[rank], output_device=rank) optimizer = AdamW(model.parameters(), lr=lr) tokenizer = BertTokenizer.from_pretrained(model_type) def update_weights(bi, di, num_batches, batch_loss): batch_loss.backward() optimizer.step() optimizer.zero_grad() if bi % 100 == 0: logger.info( f'training: device={di}; batch={bi+1}/{num_batches}; batch_error={batch_loss.item()};' ) def valid_loss_progress_log(bi, di, num_batches, batch_loss): if bi % 100 == 0: logger.info( f'validation: device={di}; batch={bi+1}/{num_batches}; val_batch_error={batch_loss.item()};' ) for i in range(epochs): model.train() train_data_loader.sampler.set_epoch(i) valid_data_loader.sampler.set_epoch(i) train_loss = run(model, train_data_loader, tokenizer, device, update_weights) if valid_data_loader is not None: with torch.no_grad(): model.eval() val_loss = run(model, valid_data_loader, tokenizer, device, valid_loss_progress_log) else: val_loss = 'N/A' logger.info( f'epoch={i}; device={rank}; train_error={train_loss}; valid_error={val_loss};' ) return model.module
def validate(student_model_without_ddp, data_loader, config, device, distributed, device_ids): teacher_model_config = config['teacher_model'] org_model, teacher_model_type = mimic_util.get_org_model( teacher_model_config, device) mimic_model = mimic_util.get_mimic_model( config, org_model, teacher_model_type, teacher_model_config, device, head_model=student_model_without_ddp) mimic_model_without_dp = mimic_model.module if isinstance( mimic_model, DataParallel) else mimic_model if distributed: mimic_model = DistributedDataParallel(mimic_model_without_dp, device_ids=device_ids) return evaluate(mimic_model, data_loader, device, split_name='Validation')
def __init__(self, seq2slate_transformer_net: Seq2SlateTransformerNet): super().__init__() self.state_dim = seq2slate_transformer_net.state_dim self.candidate_dim = seq2slate_transformer_net.candidate_dim self.num_stacked_layers = seq2slate_transformer_net.num_stacked_layers self.num_heads = seq2slate_transformer_net.num_heads self.dim_model = seq2slate_transformer_net.dim_model self.dim_feedforward = seq2slate_transformer_net.dim_feedforward self.max_src_seq_len = seq2slate_transformer_net.max_src_seq_len self.max_tgt_seq_len = seq2slate_transformer_net.max_tgt_seq_len current_device = torch.cuda.current_device() # type: ignore self.data_parallel = DistributedDataParallel( seq2slate_transformer_net.seq2slate_transformer, device_ids=[current_device], output_device=current_device, ) self.seq2slate_transformer_net = seq2slate_transformer_net
def run(args): distributed, device_ids = main_util.init_distributed_mode( args.world_size, args.dist_url) device = torch.device(args.device if torch.cuda.is_available() else 'cpu') if torch.cuda.is_available(): cudnn.benchmark = True print(args) config = yaml_util.load_yaml_file(args.config) dataset_config = config['dataset'] input_shape = config['input_shape'] train_config = config['train'] test_config = config['test'] train_loader, valid_loader, test_loader =\ dataset_util.get_data_loaders(dataset_config, batch_size=train_config['batch_size'], rough_size=train_config['rough_size'], reshape_size=input_shape[1:3], test_batch_size=test_config['batch_size'], jpeg_quality=-1, distributed=distributed) teacher_model_config = config['teacher_model'] if not args.test_only: distill(train_loader, valid_loader, input_shape, args.aux, config, device, distributed, device_ids) org_model, teacher_model_type = mimic_util.get_org_model( teacher_model_config, device) if not args.student_only: if distributed: org_model = DataParallel(org_model, device_ids=device_ids) evaluate(org_model, test_loader, device, title='[Original model]') mimic_model = mimic_util.get_mimic_model(config, org_model, teacher_model_type, teacher_model_config, device) mimic_model_without_dp = mimic_model.module if isinstance( mimic_model, DataParallel) else mimic_model file_util.save_pickle(mimic_model_without_dp, config['mimic_model']['ckpt']) if distributed: mimic_model = DistributedDataParallel(mimic_model_without_dp, device_ids=device_ids) evaluate(mimic_model, test_loader, device, title='[Mimic model]')
def __init__(self, model, training_set, batch_size, learning_rate, validation_set=None, test_set=None, checkpoint_dir=None): init_process_group('mpi') model = DistributedDataParallel(model) rank = get_rank() world_size = get_world_size() super().__init__(rank=rank, world_size=world_size, model=model, training_set=training_set, validation_set=validation_set, test_set=test_set, batch_size=batch_size, learning_rate=learning_rate, checkpoint_dir=checkpoint_dir)
def configure_ddp(self, model: LightningModule, device_ids: List[int]) -> DistributedDataParallel: """ Pass through all customizations from constructor to :class:`~torch.nn.parallel.DistributedDataParallel`. Override to define a custom DDP implementation. .. note:: This requires that your DDP implementation subclasses :class:`~torch.nn.parallel.DistributedDataParallel` and that the original LightningModule gets wrapped by :class:`~pytorch_lightning.overrides.data_parallel.LightningDistributedModule`. The default implementation is:: def configure_ddp(self, model, device_ids): model = DistributedDataParallel( LightningDistributedModule(model), device_ids=device_ids, **self._ddp_kwargs, ) return model Args: model: the LightningModule device_ids: the list of devices available Returns: the model wrapped in :class:`~torch.nn.parallel.DistributedDataParallel` """ # if unset, default `find_unused_parameters` `True` self._ddp_kwargs["find_unused_parameters"] = self._ddp_kwargs.get( "find_unused_parameters", True) model = DistributedDataParallel( module=LightningDistributedModule(model), device_ids=device_ids, **self._ddp_kwargs, ) return model
def main(): parser = get_parser() args = parser.parse_args() name = get_name(parser, args) print(f"Experiment {name}") if args.musdb is None and args.rank == 0: print( "You must provide the path to the MusDB dataset with the --musdb flag. " "To download the MusDB dataset, see https://sigsep.github.io/datasets/musdb.html.", file=sys.stderr) sys.exit(1) eval_folder = args.evals / name eval_folder.mkdir(exist_ok=True, parents=True) args.logs.mkdir(exist_ok=True) metrics_path = args.logs / f"{name}.json" eval_folder.mkdir(exist_ok=True, parents=True) args.checkpoints.mkdir(exist_ok=True, parents=True) args.models.mkdir(exist_ok=True, parents=True) if args.device is None: device = "cpu" if th.cuda.is_available(): device = "cuda" else: device = args.device th.manual_seed(args.seed) # Prevents too many threads to be started when running `museval` as it can be quite # inefficient on NUMA architectures. os.environ["OMP_NUM_THREADS"] = "1" os.environ["MKL_NUM_THREADS"] = "1" if args.world_size > 1: if device != "cuda" and args.rank == 0: print( "Error: distributed training is only available with cuda device", file=sys.stderr) sys.exit(1) th.cuda.set_device(args.rank % th.cuda.device_count()) distributed.init_process_group(backend="nccl", init_method="tcp://" + args.master, rank=args.rank, world_size=args.world_size) checkpoint = args.checkpoints / f"{name}.th" checkpoint_tmp = args.checkpoints / f"{name}.th.tmp" if args.restart and checkpoint.exists() and args.rank == 0: checkpoint.unlink() if args.test or args.test_pretrained: args.epochs = 1 args.repeat = 0 if args.test: model = load_model(args.models / args.test) else: model = load_pretrained(args.test_pretrained) elif args.tasnet: model = ConvTasNet(audio_channels=args.audio_channels, samplerate=args.samplerate, X=args.X, segment_length=4 * args.samples) else: model = Demucs( audio_channels=args.audio_channels, channels=args.channels, context=args.context, depth=args.depth, glu=args.glu, growth=args.growth, kernel_size=args.kernel_size, lstm_layers=args.lstm_layers, rescale=args.rescale, rewrite=args.rewrite, stride=args.conv_stride, resample=args.resample, samplerate=args.samplerate, segment_length=4 * args.samples, ) model.to(device) if args.init: model.load_state_dict(load_pretrained(args.init).state_dict()) if args.show: print(model) size = sizeof_fmt(4 * sum(p.numel() for p in model.parameters())) print(f"Model size {size}") return try: saved = th.load(checkpoint, map_location='cpu') except IOError: saved = SavedState() optimizer = th.optim.Adam(model.parameters(), lr=args.lr) quantizer = None quantizer = get_quantizer(model, args, optimizer) if saved.last_state is not None: model.load_state_dict(saved.last_state, strict=False) if saved.optimizer is not None: optimizer.load_state_dict(saved.optimizer) model_name = f"{name}.th" if args.save_model: if args.rank == 0: model.to("cpu") model.load_state_dict(saved.best_state) save_model(model, quantizer, args, args.models / model_name) return elif args.save_state: model_name = f"{args.save_state}.th" if args.rank == 0: model.to("cpu") model.load_state_dict(saved.best_state) state = get_state(model, quantizer) save_state(state, args.models / model_name) return if args.rank == 0: done = args.logs / f"{name}.done" if done.exists(): done.unlink() augment = [Shift(args.data_stride)] if args.augment: augment += [ FlipSign(), FlipChannels(), Scale(), Remix(group_size=args.remix_group_size) ] augment = nn.Sequential(*augment).to(device) print("Agumentation pipeline:", augment) if args.mse: criterion = nn.MSELoss() else: criterion = nn.L1Loss() # Setting number of samples so that all convolution windows are full. # Prevents hard to debug mistake with the prediction being shifted compared # to the input mixture. samples = model.valid_length(args.samples) print(f"Number of training samples adjusted to {samples}") samples = samples + args.data_stride if args.repitch: # We need a bit more audio samples, to account for potential # tempo change. samples = math.ceil(samples / (1 - 0.01 * args.max_tempo)) sources = 4 if args.raw: train_set = Rawset(args.raw / "train", samples=samples, channels=args.audio_channels, streams=range(1, sources + 1), stride=args.data_stride) valid_set = Rawset(args.raw / "valid", channels=args.audio_channels) else: if not args.metadata.is_file() and args.rank == 0: build_musdb_metadata(args.metadata, args.musdb, args.workers) if args.world_size > 1: distributed.barrier() metadata = json.load(open(args.metadata)) duration = Fraction(samples, args.samplerate) stride = Fraction(args.data_stride, args.samplerate) train_set = StemsSet(get_musdb_tracks(args.musdb, subsets=["train"], split="train"), metadata, duration=duration, stride=stride, streams=range(1, sources + 1), samplerate=args.samplerate, channels=args.audio_channels) valid_set = StemsSet(get_musdb_tracks(args.musdb, subsets=["train"], split="valid"), metadata, samplerate=args.samplerate, channels=args.audio_channels) if args.repitch: train_set = RepitchedWrapper(train_set, proba=args.repitch, max_tempo=args.max_tempo) best_loss = float("inf") for epoch, metrics in enumerate(saved.metrics): print(f"Epoch {epoch:03d}: " f"train={metrics['train']:.8f} " f"valid={metrics['valid']:.8f} " f"best={metrics['best']:.4f} " f"ms={metrics.get('true_model_size', 0):.2f}MB " f"cms={metrics.get('compressed_model_size', 0):.2f}MB " f"duration={human_seconds(metrics['duration'])}") best_loss = metrics['best'] if args.world_size > 1: dmodel = DistributedDataParallel( model, device_ids=[th.cuda.current_device()], output_device=th.cuda.current_device()) else: dmodel = model for epoch in range(len(saved.metrics), args.epochs): begin = time.time() model.train() train_loss, model_size = train_model(epoch, train_set, dmodel, criterion, optimizer, augment, quantizer=quantizer, batch_size=args.batch_size, device=device, repeat=args.repeat, seed=args.seed, diffq=args.diffq, workers=args.workers, world_size=args.world_size) model.eval() valid_loss = validate_model(epoch, valid_set, model, criterion, device=device, rank=args.rank, split=args.split_valid, overlap=args.overlap, world_size=args.world_size) ms = 0 cms = 0 if quantizer and args.rank == 0: ms = quantizer.true_model_size() cms = quantizer.compressed_model_size( num_workers=min(40, args.world_size * 10)) duration = time.time() - begin if valid_loss < best_loss and ms <= args.ms_target: best_loss = valid_loss saved.best_state = { key: value.to("cpu").clone() for key, value in model.state_dict().items() } saved.metrics.append({ "train": train_loss, "valid": valid_loss, "best": best_loss, "duration": duration, "model_size": model_size, "true_model_size": ms, "compressed_model_size": cms, }) if args.rank == 0: json.dump(saved.metrics, open(metrics_path, "w")) saved.last_state = model.state_dict() saved.optimizer = optimizer.state_dict() if args.rank == 0 and not args.test: th.save(saved, checkpoint_tmp) checkpoint_tmp.rename(checkpoint) print( f"Epoch {epoch:03d}: " f"train={train_loss:.8f} valid={valid_loss:.8f} best={best_loss:.4f} ms={ms:.2f}MB " f"cms={cms:.2f}MB " f"duration={human_seconds(duration)}") del dmodel model.load_state_dict(saved.best_state) if args.eval_cpu: device = "cpu" model.to(device) model.eval() evaluate(model, args.musdb, eval_folder, rank=args.rank, world_size=args.world_size, device=device, save=args.save, split=args.split_valid, shifts=args.shifts, overlap=args.overlap, workers=args.eval_workers) model.to("cpu") if args.rank == 0: if not (args.test or args.test_pretrained): save_model(model, quantizer, args, args.models / model_name) print("done") done.write_text("done")
def register_ddp_comm_hook( model: DistributedDataParallel, ddp_comm_state: Optional[object] = None, ddp_comm_hook: Optional[Callable] = None, ddp_comm_wrapper: Optional[Callable] = None, ) -> None: """Function to register communication hook for DDP model https://pytorch.org/docs/master/ddp_comm_hooks.html. Args: model: DDP model ddp_comm_state: state is passed to the hook and can be used to maintain and update any state information that users would like to maintain as part of the training process. Examples: error feedback in gradient compression, peers to communicate with next in GossipGrad etc. ddp_comm_hook: hook(state: object, bucket: dist._GradBucket) -> torch.futures.Future This callable function is called once the bucket is ready. The hook can perform whatever processing is needed and return a Future indicating completion of any async work (ex: allreduce). If the hook doesn't perform any communication, it can also just return a completed Future. The Future should hold the new value of grad bucket's tensors. Once a bucket is ready, c10d reducer would call this hook and use the tensors returned by the Future and copy grads to individual parameters. ddp_comm_wrapper: communication hook wrapper to support a communication hook such as FP16 compression as wrapper, which could be combined with ddp_comm_hook .. warning :: DDP communication hook needs pytorch version at least 1.8.0 .. warning :: DDP communication wrapper needs pytorch version at least 1.9.0 Post-localSGD hook needs pytorch version at least 1.9.0 Examples: >>> from torch.distributed.algorithms.ddp_comm_hooks import ( # doctest: +SKIP ... default_hooks as default, ... powerSGD_hook as powerSGD, ... post_localSGD_hook as post_localSGD, ... ) >>> >>> # fp16_compress_hook for compress gradients >>> ddp_model = ... >>> register_ddp_comm_hook( # doctest: +SKIP ... model=ddp_model, ... ddp_comm_hook=default.fp16_compress_hook, ... ) >>> >>> # powerSGD_hook >>> ddp_model = ... >>> register_ddp_comm_hook( # doctest: +SKIP ... model=ddp_model, ... ddp_comm_state=powerSGD.PowerSGDState( ... process_group=None, ... matrix_approximation_rank=1, ... start_powerSGD_iter=5000, ... ), ... ddp_comm_hook=powerSGD.powerSGD_hook, ... ) >>> >>> # post_localSGD_hook >>> subgroup, _ = torch.distributed.new_subgroups() # doctest: +SKIP >>> ddp_model = ... >>> register_ddp_comm_hook( # doctest: +SKIP ... model=ddp_model, ... state=post_localSGD.PostLocalSGDState( ... process_group=None, ... subgroup=subgroup, ... start_localSGD_iter=1_000, ... ), ... ddp_comm_hook=post_localSGD.post_localSGD_hook, ... ) >>> >>> # fp16_compress_wrapper combined with other communication hook >>> ddp_model = ... >>> register_ddp_comm_hook( # doctest: +SKIP ... model=ddp_model, ... ddp_comm_state=powerSGD.PowerSGDState( ... process_group=None, ... matrix_approximation_rank=1, ... start_powerSGD_iter=5000, ... ), ... ddp_comm_hook=powerSGD.powerSGD_hook, ... ddp_comm_wrapper=default.fp16_compress_wrapper, ... ) """ from pytorch_lightning.utilities import rank_zero_warn if not _TORCH_GREATER_EQUAL_1_8: rank_zero_warn( "Not registering DDP comm hook. To use communication hooks, please use pytorch>=1.8.0." ) return if ddp_comm_hook is None: return # inform mypy that ddp_comm_hook is callable ddp_comm_hook: Callable = ddp_comm_hook if ddp_comm_wrapper is not None: if not _TORCH_GREATER_EQUAL_1_9: rank_zero_warn( "Not applying DDP comm wrapper. To use communication wrapper, please use pytorch>=1.9.0." ) else: new_rank_zero_info( f"DDP comm wrapper is provided, apply {ddp_comm_wrapper.__qualname__}({ddp_comm_hook.__qualname__})." ) ddp_comm_hook = ddp_comm_wrapper(ddp_comm_hook) rank_zero_debug( f"Registering DDP comm hook: {ddp_comm_hook.__qualname__}.") model.register_comm_hook(state=ddp_comm_state, hook=ddp_comm_hook)
def train(cfg, args): logger = logging.getLogger('model training') train_data = build_dataloader(cfg.train_data_cfg, args.distributed) logger.info("train data: {}".format(len(train_data))) val_data = build_dataloader(cfg.val_data_cfg, args.distributed) logger.info("val data: {}".format(len(val_data))) model = MultiInstanceRecognition(cfg.model_cfg).cuda() if cfg.resume_from is not None: logger.info('loading pretrained models from {opt.continue_model}') model.load_state_dict(torch.load(cfg.resume_from)) if args.distributed: model = DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank) voc, char2id, id2char = get_vocabulary("ALLCASES_SYMBOLS") filtered_parameters = [] params_num = [] for p in filter(lambda p: p.requires_grad, model.parameters()): filtered_parameters.append(p) params_num.append(np.prod(p.size())) logger.info('Trainable params num : ', sum(params_num)) optimizer = optim.Adam(filtered_parameters, lr=cfg.lr, betas=(0.9, 0.999)) lrScheduler = lr_scheduler.MultiStepLR(optimizer, [1, 2, 3], gamma=0.1) max_iters = cfg.max_iters start_iter = 0 if cfg.resume_from is not None: start_iter = int(cfg.resume_from.split('_')[-1].split('.')[0]) logger.info('continue to train, start_iter: {start_iter}') train_data_iter = iter(train_data) val_data_iter = iter(val_data) start_time = time.time() for i in range(start_iter, max_iters): model.train() try: batch_data = next(train_data_iter) except StopIteration: train_data_iter = iter(train_data) batch_data = next(train_data_iter) data_time_s = time.time() batch_imgs, batch_imgs_path, batch_rectangles, \ batch_text_labels, batch_text_labels_mask, batch_words = \ batch_data while batch_imgs is None: batch_data = next(train_data_iter) batch_imgs, batch_imgs_path, batch_rectangles, \ batch_text_labels, batch_text_labels_mask, batch_words = \ batch_data batch_imgs = batch_imgs.cuda(non_blocking=True) batch_rectangles = batch_rectangles.cuda(non_blocking=True) batch_text_labels = batch_text_labels.cuda(non_blocking=True) data_time = time.time() - data_time_s # print(time.time() -s) # s = time.time() loss, decoder_logits = model(batch_imgs, batch_text_labels, batch_rectangles, batch_text_labels_mask) del batch_data # print(time.time() - s) # print('------') # s = time.time() loss = loss.mean() print(loss) # del loss # print(time.time() - s) # print('------') if i % cfg.train_verbose == 0: this_time = time.time() - start_time if args.distributed: loss = dist.reduce(loss, 0) log_info = "train iter :{}, time: {:.2f}, data_time: {:.2f}, Loss: {:.3f}".format( i, this_time, data_time, loss.data) logger.info(log_info) torch.cuda.empty_cache() # break optimizer.zero_grad() loss.backward() optimizer.step() del loss if i % cfg.val_iter == 0: print("--------Val iteration---------") model.eval() try: val_batch = next(val_data_iter) except StopIteration: val_data_iter = iter(val_data) val_batch = next(val_data_iter) batch_imgs, batch_imgs_path, batch_rectangles, \ batch_text_labels, batch_text_labels_mask, batch_words = \ val_batch while batch_imgs is None: val_batch = next(val_data_iter) batch_imgs, batch_imgs_path, batch_rectangles, \ batch_text_labels, batch_text_labels_mask, batch_words = \ val_batch del val_batch batch_imgs = batch_imgs.cuda(non_blocking=True) batch_rectangles = batch_rectangles.cuda(non_blocking=True) batch_text_labels = batch_text_labels.cuda(non_blocking=True) with torch.no_grad(): val_loss, val_pred_logits = model(batch_imgs, batch_text_labels, batch_rectangles, batch_text_labels_mask) pred_labels = val_pred_logits.argmax(dim=2).cpu().numpy() pred_value_str = idx2label(pred_labels, id2char, char2id) # gt_str = batch_words gt_str = [] for words in batch_words: gt_str = gt_str + words val_dec_metrics_result = calc_metrics(pred_value_str, gt_str, metrics_type="accuracy") this_time = time.time() - start_time if args.distributed: loss = dist.reduce(val_loss, 0) log_info = "val iter :{}, time: {:.2f} Loss: {:.3f}, acc: {:.2f}".format( i, this_time, loss.mean().data, val_dec_metrics_result) logger.info(log_info) del val_loss if (i + 1) % cfg.save_iter == 0: torch.save(model.state_dict(), cfg.save_name + '_{}.pth'.format(i + 1)) if i > 0 and i % cfg.lr_step == 0: # 调整学习速率 lrScheduler.step() logger.info("lr step") # torch.cuda.empty_cache() print('end the training')
def __init__( self, *, model, diffusion, data, batch_size, microbatch, lr, ema_rate, log_interval, save_interval, resume_checkpoint, use_fp16=False, fp16_scale_growth=1e-3, schedule_sampler=None, weight_decay=0.0, lr_anneal_steps=0, ): self.model = model self.diffusion = diffusion self.data = data self.batch_size = batch_size self.microbatch = microbatch if microbatch > 0 else batch_size self.lr = lr self.ema_rate = ( [ema_rate] if isinstance(ema_rate, float) else [float(x) for x in ema_rate.split(",")] ) self.log_interval = log_interval self.save_interval = save_interval self.resume_checkpoint = resume_checkpoint self.use_fp16 = use_fp16 self.fp16_scale_growth = fp16_scale_growth self.schedule_sampler = schedule_sampler or UniformSampler(diffusion) self.weight_decay = weight_decay self.lr_anneal_steps = lr_anneal_steps self.step = 0 self.resume_step = 0 self.global_batch = self.batch_size * dist.get_world_size() self.model_params = list(self.model.parameters()) self.master_params = self.model_params self.lg_loss_scale = INITIAL_LOG_LOSS_SCALE self.sync_cuda = th.cuda.is_available() self._load_and_sync_parameters() if self.use_fp16: self._setup_fp16() self.opt = AdamW(self.master_params, lr=self.lr, weight_decay=self.weight_decay) if self.resume_step: self._load_optimizer_state() # Model was resumed, either due to a restart or a checkpoint # being specified at the command line. self.ema_params = [ self._load_ema_parameters(rate) for rate in self.ema_rate ] else: self.ema_params = [ copy.deepcopy(self.master_params) for _ in range(len(self.ema_rate)) ] if th.cuda.is_available(): self.use_ddp = True self.ddp_model = DDP( self.model, device_ids=[dist_util.dev()], output_device=dist_util.dev(), broadcast_buffers=False, bucket_cap_mb=128, find_unused_parameters=False, ) else: if dist.get_world_size() > 1: logger.warn( "Distributed training requires CUDA. " "Gradients will not be synchronized properly!" ) self.use_ddp = False self.ddp_model = self.model
class TrainLoop: def __init__( self, *, model, diffusion, data, batch_size, microbatch, lr, ema_rate, log_interval, save_interval, resume_checkpoint, use_fp16=False, fp16_scale_growth=1e-3, schedule_sampler=None, weight_decay=0.0, lr_anneal_steps=0, ): self.model = model self.diffusion = diffusion self.data = data self.batch_size = batch_size self.microbatch = microbatch if microbatch > 0 else batch_size self.lr = lr self.ema_rate = ( [ema_rate] if isinstance(ema_rate, float) else [float(x) for x in ema_rate.split(",")] ) self.log_interval = log_interval self.save_interval = save_interval self.resume_checkpoint = resume_checkpoint self.use_fp16 = use_fp16 self.fp16_scale_growth = fp16_scale_growth self.schedule_sampler = schedule_sampler or UniformSampler(diffusion) self.weight_decay = weight_decay self.lr_anneal_steps = lr_anneal_steps self.step = 0 self.resume_step = 0 self.global_batch = self.batch_size * dist.get_world_size() self.model_params = list(self.model.parameters()) self.master_params = self.model_params self.lg_loss_scale = INITIAL_LOG_LOSS_SCALE self.sync_cuda = th.cuda.is_available() self._load_and_sync_parameters() if self.use_fp16: self._setup_fp16() self.opt = AdamW(self.master_params, lr=self.lr, weight_decay=self.weight_decay) if self.resume_step: self._load_optimizer_state() # Model was resumed, either due to a restart or a checkpoint # being specified at the command line. self.ema_params = [ self._load_ema_parameters(rate) for rate in self.ema_rate ] else: self.ema_params = [ copy.deepcopy(self.master_params) for _ in range(len(self.ema_rate)) ] if th.cuda.is_available(): self.use_ddp = True self.ddp_model = DDP( self.model, device_ids=[dist_util.dev()], output_device=dist_util.dev(), broadcast_buffers=False, bucket_cap_mb=128, find_unused_parameters=False, ) else: if dist.get_world_size() > 1: logger.warn( "Distributed training requires CUDA. " "Gradients will not be synchronized properly!" ) self.use_ddp = False self.ddp_model = self.model def _load_and_sync_parameters(self): resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint if resume_checkpoint: self.resume_step = parse_resume_step_from_filename(resume_checkpoint) if dist.get_rank() == 0: logger.log(f"loading model from checkpoint: {resume_checkpoint}...") self.model.load_state_dict( dist_util.load_state_dict( resume_checkpoint, map_location=dist_util.dev() ) ) dist_util.sync_params(self.model.parameters()) def _load_ema_parameters(self, rate): ema_params = copy.deepcopy(self.master_params) main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint ema_checkpoint = find_ema_checkpoint(main_checkpoint, self.resume_step, rate) if ema_checkpoint: if dist.get_rank() == 0: logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...") state_dict = dist_util.load_state_dict( ema_checkpoint, map_location=dist_util.dev() ) ema_params = self._state_dict_to_master_params(state_dict) dist_util.sync_params(ema_params) return ema_params def _load_optimizer_state(self): main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint opt_checkpoint = bf.join( bf.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt" ) if bf.exists(opt_checkpoint): logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}") state_dict = dist_util.load_state_dict( opt_checkpoint, map_location=dist_util.dev() ) self.opt.load_state_dict(state_dict) def _setup_fp16(self): self.master_params = make_master_params(self.model_params) self.model.convert_to_fp16() def run_loop(self): while ( not self.lr_anneal_steps or self.step + self.resume_step < self.lr_anneal_steps ): batch, cond = next(self.data) self.run_step(batch, cond) if self.step % self.log_interval == 0: logger.dumpkvs() if self.step % self.save_interval == 0: self.save() # Run for a finite amount of time in integration tests. if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0: return self.step += 1 # Save the last checkpoint if it wasn't already saved. if (self.step - 1) % self.save_interval != 0: self.save() def run_step(self, batch, cond): self.forward_backward(batch, cond) if self.use_fp16: self.optimize_fp16() else: self.optimize_normal() self.log_step() def forward_backward(self, batch, cond): zero_grad(self.model_params) for i in range(0, batch.shape[0], self.microbatch): micro = batch[i : i + self.microbatch].to(dist_util.dev()) micro_cond = { k: v[i : i + self.microbatch].to(dist_util.dev()) for k, v in cond.items() } last_batch = (i + self.microbatch) >= batch.shape[0] t, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev()) compute_losses = functools.partial( self.diffusion.training_losses, self.ddp_model, micro, t, model_kwargs=micro_cond, ) if last_batch or not self.use_ddp: losses = compute_losses() else: with self.ddp_model.no_sync(): losses = compute_losses() if isinstance(self.schedule_sampler, LossAwareSampler): self.schedule_sampler.update_with_local_losses( t, losses["loss"].detach() ) loss = (losses["loss"] * weights).mean() log_loss_dict( self.diffusion, t, {k: v * weights for k, v in losses.items()} ) if self.use_fp16: loss_scale = 2 ** self.lg_loss_scale (loss * loss_scale).backward() else: loss.backward() def optimize_fp16(self): if any(not th.isfinite(p.grad).all() for p in self.model_params): self.lg_loss_scale -= 1 logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}") return model_grads_to_master_grads(self.model_params, self.master_params) self.master_params[0].grad.mul_(1.0 / (2 ** self.lg_loss_scale)) self._log_grad_norm() self._anneal_lr() self.opt.step() for rate, params in zip(self.ema_rate, self.ema_params): update_ema(params, self.master_params, rate=rate) master_params_to_model_params(self.model_params, self.master_params) self.lg_loss_scale += self.fp16_scale_growth def optimize_normal(self): self._log_grad_norm() self._anneal_lr() self.opt.step() for rate, params in zip(self.ema_rate, self.ema_params): update_ema(params, self.master_params, rate=rate) def _log_grad_norm(self): sqsum = 0.0 for p in self.master_params: sqsum += (p.grad ** 2).sum().item() logger.logkv_mean("grad_norm", np.sqrt(sqsum)) def _anneal_lr(self): if not self.lr_anneal_steps: return frac_done = (self.step + self.resume_step) / self.lr_anneal_steps lr = self.lr * (1 - frac_done) for param_group in self.opt.param_groups: param_group["lr"] = lr def log_step(self): logger.logkv("step", self.step + self.resume_step) logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch) if self.use_fp16: logger.logkv("lg_loss_scale", self.lg_loss_scale) def save(self): def save_checkpoint(rate, params): state_dict = self._master_params_to_state_dict(params) if dist.get_rank() == 0: logger.log(f"saving model {rate}...") if not rate: filename = f"model{(self.step+self.resume_step):06d}.pt" else: filename = f"ema_{rate}_{(self.step+self.resume_step):06d}.pt" with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f: th.save(state_dict, f) save_checkpoint(0, self.master_params) for rate, params in zip(self.ema_rate, self.ema_params): save_checkpoint(rate, params) if dist.get_rank() == 0: with bf.BlobFile( bf.join(get_blob_logdir(), f"opt{(self.step+self.resume_step):06d}.pt"), "wb", ) as f: th.save(self.opt.state_dict(), f) dist.barrier() def _master_params_to_state_dict(self, master_params): if self.use_fp16: master_params = unflatten_master_params( self.model.parameters(), master_params ) state_dict = self.model.state_dict() for i, (name, _value) in enumerate(self.model.named_parameters()): assert name in state_dict state_dict[name] = master_params[i] return state_dict def _state_dict_to_master_params(self, state_dict): params = [state_dict[name] for name, _ in self.model.named_parameters()] if self.use_fp16: return make_master_params(params) else: return params
def distill(train_loader, valid_loader, input_shape, aux_weight, config, device, distributed, device_ids): teacher_model_config = config['teacher_model'] teacher_model, teacher_model_type = mimic_util.get_teacher_model( teacher_model_config, input_shape, device) module_util.freeze_module_params(teacher_model) student_model_config = config['student_model'] student_model = mimic_util.get_student_model(teacher_model_type, student_model_config, config['dataset']['name']) student_model = student_model.to(device) start_epoch, best_valid_acc = mimic_util.resume_from_ckpt( student_model_config['ckpt'], student_model, is_student=True) if best_valid_acc is None: best_valid_acc = 0.0 train_config = config['train'] criterion_config = train_config['criterion'] criterion = func_util.get_loss(criterion_config['type'], criterion_config['params']) optim_config = train_config['optimizer'] optimizer = func_util.get_optimizer(student_model, optim_config['type'], optim_config['params']) scheduler_config = train_config['scheduler'] scheduler = func_util.get_scheduler(optimizer, scheduler_config['type'], scheduler_config['params']) interval = train_config['interval'] if interval <= 0: num_batches = len(train_loader) interval = num_batches // 20 if num_batches >= 20 else 1 student_model_without_ddp = student_model if distributed: teacher_model = DataParallel(teacher_model, device_ids=device_ids) student_model = DistributedDataParallel(student_model, device_ids=device_ids) student_model_without_ddp = student_model.module ckpt_file_path = student_model_config['ckpt'] end_epoch = start_epoch + train_config['epoch'] start_time = time.time() for epoch in range(start_epoch, end_epoch): if distributed: train_loader.sampler.set_epoch(epoch) distill_one_epoch(student_model, teacher_model, train_loader, optimizer, criterion, epoch, device, interval, aux_weight) valid_acc = validate(student_model, valid_loader, config, device, distributed, device_ids) if valid_acc > best_valid_acc and main_util.is_main_process(): print( 'Updating ckpt (Best top1 accuracy: {:.4f} -> {:.4f})'.format( best_valid_acc, valid_acc)) best_valid_acc = valid_acc save_ckpt(student_model_without_ddp, epoch, best_valid_acc, ckpt_file_path, teacher_model_type) scheduler.step() dist.barrier() total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str)) del teacher_model del student_model
# # 数值归一化 # cam -= np.min(cam) # cam /= np.max(cam) # # resize to 224*224 # cam = cv2.resize(cam, (224, 224)) # return cam print("gradient: ", gradient.shape) print("feature: ", feature.shape) os.environ['CUDA_VISIBLE_DEVICES'] = "1,3" torch.distributed.init_process_group(backend="nccl", init_method='tcp://localhost:12556', rank=0, world_size=1) model = ThinResNet() model = model.cuda() model = DistributedDataParallel(model) gc = GradCAM(model, 'layer4') x = torch.randn((20, 1, 224, 224)).cuda() l = torch.range(0, 19).long().unsqueeze(1).cuda() y = model(x) # cam = gc(x, l) print(cam.shape)
class BaseTrainer: def __init__( self, task, model, dataset, optimizer, identifier, run_dir=None, is_debug=False, is_vis=False, print_every=100, seed=None, logger="tensorboard", local_rank=0, amp=False, name="base_trainer", ): self.name = name if torch.cuda.is_available(): self.device = local_rank else: self.device = "cpu" if run_dir is None: run_dir = os.getcwd() run_dir = Path(run_dir) timestamp = torch.tensor(datetime.datetime.now().timestamp()).to( self.device) # create directories from master rank only distutils.broadcast(timestamp, 0) timestamp = datetime.datetime.fromtimestamp(timestamp).strftime( "%Y-%m-%d-%H-%M-%S") if identifier: timestamp += "-{}".format(identifier) self.config = { "task": task, "model": model.pop("name"), "model_attributes": model, "optim": optimizer, "logger": logger, "amp": amp, "cmd": { "identifier": identifier, "print_every": print_every, "seed": seed, "timestamp": timestamp, "checkpoint_dir": str(run_dir / "checkpoints" / timestamp), "results_dir": str(run_dir / "results" / timestamp), "logs_dir": str(run_dir / "logs" / logger / timestamp), }, } # AMP Scaler self.scaler = torch.cuda.amp.GradScaler() if amp else None if isinstance(dataset, list): self.config["dataset"] = dataset[0] if len(dataset) > 1: self.config["val_dataset"] = dataset[1] if len(dataset) > 2: self.config["test_dataset"] = dataset[2] else: self.config["dataset"] = dataset if not is_debug and distutils.is_master(): os.makedirs(self.config["cmd"]["checkpoint_dir"], exist_ok=True) os.makedirs(self.config["cmd"]["results_dir"], exist_ok=True) os.makedirs(self.config["cmd"]["logs_dir"], exist_ok=True) self.is_debug = is_debug self.is_vis = is_vis if distutils.is_master(): print(yaml.dump(self.config, default_flow_style=False)) self.load() self.evaluator = Evaluator(task=name) def load(self): self.load_seed_from_config() self.load_logger() self.load_task() self.load_model() self.load_criterion() self.load_optimizer() self.load_extras() # Note: this function is now deprecated. We build config outside of trainer. # See build_config in ocpmodels.common.utils.py. def load_config_from_yaml_and_cmd(self, args): self.config = build_config(args) # AMP Scaler self.scaler = (torch.cuda.amp.GradScaler() if self.config["amp"] else None) # device self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu") # Are we just running sanity checks? self.is_debug = args.debug self.is_vis = args.vis # timestamps and directories args.timestamp = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S") if args.identifier: args.timestamp += "-{}".format(args.identifier) args.checkpoint_dir = os.path.join("checkpoints", args.timestamp) args.results_dir = os.path.join("results", args.timestamp) args.logs_dir = os.path.join("logs", self.config["logger"], args.timestamp) print(yaml.dump(self.config, default_flow_style=False)) for arg in vars(args): print("{:<20}: {}".format(arg, getattr(args, arg))) # TODO(abhshkdz): Handle these parameters better. Maybe move to yaml. self.config["cmd"] = args.__dict__ del args if not self.is_debug: os.makedirs(self.config["cmd"]["checkpoint_dir"], exist_ok=True) os.makedirs(self.config["cmd"]["results_dir"], exist_ok=True) os.makedirs(self.config["cmd"]["logs_dir"], exist_ok=True) # Dump config parameters json.dump( self.config, open( os.path.join(self.config["cmd"]["checkpoint_dir"], "config.json"), "w", ), ) def load_seed_from_config(self): # https://pytorch.org/docs/stable/notes/randomness.html seed = self.config["cmd"]["seed"] if seed is None: return random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False def load_logger(self): self.logger = None if not self.is_debug and distutils.is_master(): assert (self.config["logger"] is not None), "Specify logger in config" self.logger = registry.get_logger_class(self.config["logger"])( self.config) def load_task(self): print("### Loading dataset: {}".format(self.config["task"]["dataset"])) dataset = registry.get_dataset_class(self.config["task"]["dataset"])( self.config["dataset"]) if self.config["task"]["dataset"] in ["qm9", "dogss"]: num_targets = dataset.data.y.shape[-1] if ("label_index" in self.config["task"] and self.config["task"]["label_index"] is not False): dataset.data.y = dataset.data.y[:, int(self.config["task"] ["label_index"])] num_targets = 1 else: num_targets = 1 self.num_targets = num_targets ( self.train_loader, self.val_loader, self.test_loader, ) = dataset.get_dataloaders( batch_size=int(self.config["optim"]["batch_size"])) # Normalizer for the dataset. # Compute mean, std of training set labels. self.normalizers = {} if self.config["dataset"].get("normalize_labels", True): self.normalizers["target"] = Normalizer( self.train_loader.dataset.data.y[ self.train_loader.dataset.__indices__], self.device, ) # If we're computing gradients wrt input, set mean of normalizer to 0 -- # since it is lost when compute dy / dx -- and std to forward target std if "grad_input" in self.config["task"]: if self.config["dataset"].get("normalize_labels", True): self.normalizers["grad_target"] = Normalizer( self.train_loader.dataset.data.y[ self.train_loader.dataset.__indices__], self.device, ) self.normalizers["grad_target"].mean.fill_(0) if self.is_vis and self.config["task"]["dataset"] != "qm9": # Plot label distribution. plots = [ plot_histogram( self.train_loader.dataset.data.y.tolist(), xlabel="{}/raw".format(self.config["task"]["labels"][0]), ylabel="# Examples", title="Split: train", ), plot_histogram( self.val_loader.dataset.data.y.tolist(), xlabel="{}/raw".format(self.config["task"]["labels"][0]), ylabel="# Examples", title="Split: val", ), plot_histogram( self.test_loader.dataset.data.y.tolist(), xlabel="{}/raw".format(self.config["task"]["labels"][0]), ylabel="# Examples", title="Split: test", ), ] self.logger.log_plots(plots) def load_model(self): # Build model if distutils.is_master(): print("### Loading model: {}".format(self.config["model"])) # TODO(abhshkdz): Eventually move towards computing features on-the-fly # and remove dependence from `.edge_attr`. bond_feat_dim = None if self.config["task"]["dataset"] in [ "trajectory_lmdb", "single_point_lmdb", ]: bond_feat_dim = self.config["model_attributes"].get( "num_gaussians", 50) else: raise NotImplementedError self.model = registry.get_model_class(self.config["model"])( self.train_loader.dataset[0].x.shape[-1] if hasattr(self.train_loader.dataset[0], "x") and self.train_loader.dataset[0].x is not None else None, bond_feat_dim, self.num_targets, **self.config["model_attributes"], ).to(self.device) if distutils.is_master(): print("### Loaded {} with {} parameters.".format( self.model.__class__.__name__, self.model.num_params)) if self.logger is not None: self.logger.watch(self.model) self.model = OCPDataParallel( self.model, output_device=self.device, num_gpus=1, ) if distutils.initialized(): self.model = DistributedDataParallel(self.model, device_ids=[self.device]) def load_pretrained(self, checkpoint_path=None, ddp_to_dp=False): if checkpoint_path is None or os.path.isfile(checkpoint_path) is False: print(f"Checkpoint: {checkpoint_path} not found!") return False print("### Loading checkpoint from: {}".format(checkpoint_path)) checkpoint = torch.load(checkpoint_path) # Load model, optimizer, normalizer state dict. # if trained with ddp and want to load in non-ddp, modify keys from # module.module.. -> module.. if ddp_to_dp: new_dict = OrderedDict() for k, v in checkpoint["state_dict"].items(): name = k[7:] new_dict[name] = v self.model.load_state_dict(new_dict) else: self.model.load_state_dict(checkpoint["state_dict"]) self.optimizer.load_state_dict(checkpoint["optimizer"]) for key in checkpoint["normalizers"]: if key in self.normalizers: self.normalizers[key].load_state_dict( checkpoint["normalizers"][key]) if self.scaler and checkpoint["amp"]: self.scaler.load_state_dict(checkpoint["amp"]) return True # TODO(abhshkdz): Rename function to something nicer. # TODO(abhshkdz): Support multiple loss functions. def load_criterion(self): self.criterion = self.config["optim"].get("criterion", nn.L1Loss()) def load_optimizer(self): self.optimizer = optim.AdamW( self.model.parameters(), self.config["optim"]["lr_initial"], # weight_decay=3.0 ) def load_extras(self): # learning rate scheduler. scheduler_lambda_fn = lambda x: warmup_lr_lambda( x, self.config["optim"]) self.scheduler = optim.lr_scheduler.LambdaLR( self.optimizer, lr_lambda=scheduler_lambda_fn) # metrics. self.meter = Meter(split="train") def save(self, epoch, metrics): if not self.is_debug and distutils.is_master(): save_checkpoint( { "epoch": epoch, "state_dict": self.model.state_dict(), "optimizer": self.optimizer.state_dict(), "normalizers": { key: value.state_dict() for key, value in self.normalizers.items() }, "config": self.config, "val_metrics": metrics, "amp": self.scaler.state_dict() if self.scaler else None, }, self.config["cmd"]["checkpoint_dir"], ) def train(self, max_epochs=None, return_metrics=False): # TODO(abhshkdz): Timers for dataloading and forward pass. num_epochs = (max_epochs if max_epochs is not None else self.config["optim"]["max_epochs"]) for epoch in range(num_epochs): self.model.train() for i, batch in enumerate(self.train_loader): batch = batch.to(self.device) # Forward, loss, backward. out, metrics = self._forward(batch) loss = self._compute_loss(out, batch) self._backward(loss) # Update meter. meter_update_dict = { "epoch": epoch + (i + 1) / len(self.train_loader), "loss": loss.item(), } meter_update_dict.update(metrics) self.meter.update(meter_update_dict) # Make plots. if self.logger is not None: self.logger.log( meter_update_dict, step=epoch * len(self.train_loader) + i + 1, split="train", ) # Print metrics. if i % self.config["cmd"]["print_every"] == 0: print(self.meter) self.scheduler.step() with torch.no_grad(): if self.val_loader is not None: v_loss, v_mae = self.validate(split="val", epoch=epoch) if self.test_loader is not None: test_loss, test_mae = self.validate(split="test", epoch=epoch) if not self.is_debug: save_checkpoint( { "epoch": epoch + 1, "state_dict": self.model.state_dict(), "optimizer": self.optimizer.state_dict(), "normalizers": { key: value.state_dict() for key, value in self.normalizers.items() }, "config": self.config, "amp": self.scaler.state_dict() if self.scaler else None, }, self.config["cmd"]["checkpoint_dir"], ) if return_metrics: return { "training_loss": float(self.meter.loss.global_avg), "training_mae": float(self.meter.meters[ self.config["task"]["labels"][0] + "/" + self.config["task"]["metric"]].global_avg), "validation_loss": v_loss, "validation_mae": v_mae, "test_loss": test_loss, "test_mae": test_mae, } def validate(self, split="val", epoch=None): if distutils.is_master(): print("### Evaluating on {}.".format(split)) self.model.eval() evaluator, metrics = Evaluator(task=self.name), {} rank = distutils.get_rank() loader = self.val_loader if split == "val" else self.test_loader for i, batch in tqdm( enumerate(loader), total=len(loader), position=rank, desc="device {}".format(rank), ): # Forward. with torch.cuda.amp.autocast(enabled=self.scaler is not None): out = self._forward(batch) loss = self._compute_loss(out, batch) # Compute metrics. metrics = self._compute_metrics(out, batch, evaluator, metrics) metrics = evaluator.update("loss", loss.item(), metrics) aggregated_metrics = {} for k in metrics: aggregated_metrics[k] = { "total": distutils.all_reduce(metrics[k]["total"], average=False, device=self.device), "numel": distutils.all_reduce(metrics[k]["numel"], average=False, device=self.device), } aggregated_metrics[k]["metric"] = (aggregated_metrics[k]["total"] / aggregated_metrics[k]["numel"]) metrics = aggregated_metrics log_dict = {k: metrics[k]["metric"] for k in metrics} log_dict.update({"epoch": epoch + 1}) if distutils.is_master(): log_str = ["{}: {:.4f}".format(k, v) for k, v in log_dict.items()] print(", ".join(log_str)) # Make plots. if self.logger is not None and epoch is not None: self.logger.log( log_dict, step=(epoch + 1) * len(self.train_loader), split=split, ) return metrics def _forward(self, batch, compute_metrics=True): out = {} # enable gradient wrt input. if "grad_input" in self.config["task"]: inp_for_grad = batch.pos batch.pos = batch.pos.requires_grad_(True) # forward pass. if self.config["model_attributes"].get("regress_forces", False): output, output_forces = self.model(batch) else: output = self.model(batch) if output.shape[-1] == 1: output = output.view(-1) out["output"] = output force_output = None if self.config["model_attributes"].get("regress_forces", False): out["force_output"] = output_forces force_output = output_forces if ("grad_input" in self.config["task"] and self.config["model_attributes"].get( "regress_forces", False) is False): force_output = -1 * torch.autograd.grad( output, inp_for_grad, grad_outputs=torch.ones_like(output), create_graph=True, retain_graph=True, )[0] out["force_output"] = force_output if not compute_metrics: return out, None metrics = {} if self.config["dataset"].get("normalize_labels", True): errors = eval(self.config["task"]["metric"])( self.normalizers["target"].denorm(output).cpu(), batch.y.cpu()).view(-1) else: errors = eval(self.config["task"]["metric"])( output.cpu(), batch.y.cpu()).view(-1) if ("label_index" in self.config["task"] and self.config["task"]["label_index"] is not False): # TODO(abhshkdz): Get rid of this edge case for QM9. # This is only because QM9 has multiple targets and we can either # jointly predict all of them or one particular target. metrics["{}/{}".format( self.config["task"]["labels"][self.config["task"] ["label_index"]], self.config["task"]["metric"], )] = errors[0] else: for i, label in enumerate(self.config["task"]["labels"]): metrics["{}/{}".format( label, self.config["task"]["metric"])] = errors[i] if "grad_input" in self.config["task"]: force_pred = force_output force_target = batch.force if self.config["task"].get("eval_on_free_atoms", True): mask = batch.fixed == 0 force_pred = force_pred[mask] force_target = force_target[mask] if self.config["dataset"].get("normalize_labels", True): grad_input_errors = eval(self.config["task"]["metric"])( self.normalizers["grad_target"].denorm(force_pred).cpu(), force_target.cpu(), ) else: grad_input_errors = eval(self.config["task"]["metric"])( force_pred.cpu(), force_target.cpu()) metrics["force_x/{}".format( self.config["task"]["metric"])] = grad_input_errors[0] metrics["force_y/{}".format( self.config["task"]["metric"])] = grad_input_errors[1] metrics["force_z/{}".format( self.config["task"]["metric"])] = grad_input_errors[2] return out, metrics def _compute_loss(self, out, batch): loss = [] if self.config["dataset"].get("normalize_labels", True): target_normed = self.normalizers["target"].norm(batch.y) else: target_normed = batch.y loss.append(self.criterion(out["output"], target_normed)) # TODO(abhshkdz): Test support for gradients wrt input. # TODO(abhshkdz): Make this general; remove dependence on `.forces`. if "grad_input" in self.config["task"]: if self.config["dataset"].get("normalize_labels", True): grad_target_normed = self.normalizers["grad_target"].norm( batch.force) else: grad_target_normed = batch.force # Force coefficient = 30 has been working well for us. force_mult = self.config["optim"].get("force_coefficient", 30) if self.config["task"].get("train_on_free_atoms", False): mask = batch.fixed == 0 loss.append(force_mult * self.criterion( out["force_output"][mask], grad_target_normed[mask])) else: loss.append( force_mult * self.criterion(out["force_output"], grad_target_normed)) # Sanity check to make sure the compute graph is correct. for lc in loss: assert hasattr(lc, "grad_fn") loss = sum(loss) return loss def _backward(self, loss): self.optimizer.zero_grad() loss.backward() # TODO(abhshkdz): Add support for gradient clipping. if self.scaler: self.scaler.step(self.optimizer) self.scaler.update() else: self.optimizer.step() def save_results(self, predictions, results_file, keys): if results_file is None: return results_file_path = os.path.join( self.config["cmd"]["results_dir"], f"{self.name}_{results_file}_{distutils.get_rank()}.npz", ) np.savez_compressed( results_file_path, ids=predictions["id"], **{key: predictions[key] for key in keys}, ) distutils.synchronize() if distutils.is_master(): gather_results = defaultdict(list) full_path = os.path.join( self.config["cmd"]["results_dir"], f"{self.name}_{results_file}.npz", ) for i in range(distutils.get_world_size()): rank_path = os.path.join( self.config["cmd"]["results_dir"], f"{self.name}_{results_file}_{i}.npz", ) rank_results = np.load(rank_path, allow_pickle=True) gather_results["ids"].extend(rank_results["ids"]) for key in keys: gather_results[key].extend(rank_results[key]) os.remove(rank_path) # Because of how distributed sampler works, some system ids # might be repeated to make no. of samples even across GPUs. _, idx = np.unique(gather_results["ids"], return_index=True) gather_results["ids"] = np.array(gather_results["ids"])[idx] for k in keys: if k == "forces": gather_results[k] = np.array(gather_results[k], dtype=object)[idx] else: gather_results[k] = np.array(gather_results[k])[idx] print(f"Writing results to {full_path}") np.savez_compressed(full_path, **gather_results)