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 device.type == 'cuda': cudnn.benchmark = True print(args) config = yaml_util.load_yaml_file(args.config) input_shape = config['input_shape'] ckpt_file_path = config['autoencoder']['ckpt'] train_loader, valid_loader, test_loader = main_util.get_data_loaders( config, distributed) if not args.test_only: train(train_loader, valid_loader, input_shape, config, device, distributed, device_ids) autoencoder, _ = ae_util.get_autoencoder(config, device) resume_from_ckpt(ckpt_file_path, autoencoder) extended_model, model = ae_util.get_extended_model(autoencoder, config, input_shape, device) if not args.extended_only: if device.type == 'cuda': model = DistributedDataParallel(model, device_ids=device_ids) if distributed \ else DataParallel(model) evaluate(model, test_loader, device, title='[Original model]') if device.type == 'cuda': extended_model = DistributedDataParallel(extended_model, device_ids=device_ids) if distributed \ else DataParallel(extended_model) evaluate(extended_model, test_loader, device, title='[Mimic model]')
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 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) freeze_modules(student_model, student_model_config) 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, device, args.skip_teacher_eval, args.transform_bottleneck)
def configure_ddp(self): self.pre_configure_ddp() self._model = DistributedDataParallel( LightningDistributedModule(self.model), **self._ddp_kwargs, ) self._register_ddp_hooks()
def configure_ddp(self): self.pre_configure_ddp() self._model = DistributedDataParallel( LightningDistributedModule(self.model), device_ids=self.determine_ddp_device_ids(), **self._ddp_kwargs, )
def load_vaes(H, logprint): vae = VAE(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(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) vae = DistributedDataParallel(vae, device_ids=[H.local_rank], output_device=H.local_rank) 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 train(train_loader, valid_loader, input_shape, config, device, distributed, device_ids): ae_without_ddp, ae_type = ae_util.get_autoencoder(config, device) head_model = ae_util.get_head_model(config, input_shape, device) module_util.freeze_module_params(head_model) ckpt_file_path = config['autoencoder']['ckpt'] start_epoch, best_valid_acc = resume_from_ckpt(ckpt_file_path, ae_without_ddp) 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(ae_without_ddp, 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 autoencoder = ae_without_ddp if distributed: autoencoder = DistributedDataParallel(ae_without_ddp, device_ids=device_ids) head_model = DataParallel(head_model, device_ids=device_ids) elif device.type == 'cuda': autoencoder = DataParallel(ae_without_ddp) head_model = DataParallel(head_model) 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) train_epoch(autoencoder, head_model, train_loader, optimizer, criterion, epoch, device, interval) valid_acc = validate(ae_without_ddp, 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(ae_without_ddp, epoch, best_valid_acc, ckpt_file_path, ae_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 head_model
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 __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 _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, 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 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 load_model(self): super(EnergyTrainer, self).load_model() self.model = OCPDataParallel( self.model, output_device=self.device, num_gpus=self.config["optim"].get("num_gpus", 1), ) if distutils.initialized(): self.model = DistributedDataParallel(self.model, device_ids=[self.device])
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): super(DistributedEnergyTrainer, self).load_model() self.model = OCPDataParallel( self.model, output_device=self.device, num_gpus=self.config["optim"].get("num_gpus", 1), ) self.model = DistributedDataParallel(self.model, device_ids=[self.device], find_unused_parameters=True)
def load_model(self): super(DistributedForcesTrainer, self).load_model() self.model = OCPDataParallel( self.model, output_device=self.device, num_gpus=1, ) self.model = DistributedDataParallel(self.model, device_ids=[self.device], find_unused_parameters=True)
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, 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 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 load_model(self): # Build model if distutils.is_master(): logging.info(f"Loading model: {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 loader = self.train_loader or self.val_loader or self.test_loader self.model = registry.get_model_class(self.config["model"])( loader.dataset[0].x.shape[-1] if loader and hasattr(loader.dataset[0], "x") and 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(): logging.info( f"Loaded {self.model.__class__.__name__} with " f"{self.model.num_params} parameters." ) if self.logger is not None: self.logger.watch(self.model) self.model = OCPDataParallel( self.model, output_device=self.device, num_gpus=1 if not self.cpu else 0, ) if distutils.initialized(): self.model = DistributedDataParallel( self.model, device_ids=[self.device] )
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 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, 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, 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 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')