def main(args): print(args) if not torch.cuda.is_available(): raise NotImplementedError('Training on CPU is not supported') torch.cuda.set_device(args.device_id) torch.manual_seed(args.seed) # Load dataset splits = ['train', 'valid'] if data.has_binary_files(args.data, splits): dataset = data.load_dataset( args.data, splits, args.source_lang, args.target_lang) else: dataset = data.load_raw_text_dataset( args.data, splits, args.source_lang, args.target_lang) if args.source_lang is None or args.target_lang is None: # record inferred languages in args, so that it's saved in checkpoints args.source_lang, args.target_lang = dataset.src, dataset.dst print('| [{}] dictionary: {} types'.format(dataset.src, len(dataset.src_dict))) print('| [{}] dictionary: {} types'.format(dataset.dst, len(dataset.dst_dict))) for split in splits: print('| {} {} {} examples'.format(args.data, split, len(dataset.splits[split]))) # Build model and criterion model = models.build_model(args, dataset.src_dict, dataset.dst_dict) criterion = criterions.build_criterion(args, dataset.src_dict, dataset.dst_dict) print('| model {}, criterion {}'.format(args.arch, criterion.__class__.__name__)) print('| num. model params: {}'.format(sum(p.data.numel() for p in model.parameters()))) # Build trainer trainer = Trainer(args, model, criterion) print('| training on {} GPUs'.format(args.distributed_world_size)) print('| max tokens per GPU = {} and max sentences per GPU = {}'.format( args.max_tokens, args.max_sentences, )) # Load the latest checkpoint if one is available os.makedirs(args.save_dir, exist_ok=True) checkpoint_path = os.path.join(args.save_dir, args.restore_file) extra_state = trainer.load_checkpoint(checkpoint_path) if extra_state is not None: epoch = extra_state['epoch'] batch_offset = extra_state['batch_offset'] print('| loaded checkpoint {} (epoch {})'.format(checkpoint_path, epoch)) if batch_offset == 0: trainer.lr_step(epoch) epoch += 1 else: epoch, batch_offset = 1, 0 # Train until the learning rate gets too small max_epoch = args.max_epoch or math.inf lr = trainer.get_lr() train_meter = StopwatchMeter() train_meter.start() while lr > args.min_lr and epoch <= max_epoch: # train for one epoch train(args, trainer, dataset, epoch, batch_offset) # evaluate on validate set for k, subset in enumerate(args.valid_subset.split(',')): val_loss = validate(args, trainer, dataset, subset, epoch) if k == 0: # only use first validation loss to update the learning schedule lr = trainer.lr_step(epoch, val_loss) # save checkpoint if not args.no_save: save_checkpoint(trainer, args, epoch, 0, val_loss) epoch += 1 batch_offset = 0 train_meter.stop() print('| done training in {:.1f} seconds'.format(train_meter.sum))
losses = -np.ones(shape=(len(xcoordinates), len(ycoordinates))) accuracies = -np.ones(shape=(len(xcoordinates), len(ycoordinates))) f['train_loss'] = losses f['train_acc'] = accuracies inds, coords, inds_nums = scheduler.get_job_indices( losses, xcoordinates, ycoordinates, None) start_time = time.time() total_sync = 0.0 xtra_state = trainer.load_checkpoint( os.path.join(args_transformer.save_dir, args_transformer.restore_file), args_transformer.reset_optimizer, args_transformer.reset_lr_scheduler, eval(args_transformer.optimizer_overrides), reset_meters=args_transformer.reset_meters, ) extra_state, epoch_itr = checkpoint_utils.load_checkpoint( args_transformer, trainer) update_freq = args_transformer.update_freq[epoch_itr.epoch - 1] if epoch_itr.epoch <= len(args_transformer.update_freq) else \ args_transformer.update_freq[-1] # Initialize data iterator itr = epoch_itr.next_epoch_itr( fix_batches_to_gpus=args_transformer.fix_batches_to_gpus, shuffle=(epoch_itr.epoch >= args_transformer.curriculum), )
def main(args, init_distributed=False): utils.import_user_module(args) # Initialize CUDA and distributed training random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) # Setup task, (should be default, translation) task = tasks.setup_task(args) # Build model and criterion model = task.build_model(args) criterion = task.build_criterion(args) # Build trainer trainer = Trainer(args, task, model, criterion) initial_state_checkpoint = str(pathlib.Path(args.save_dir) / 'initial.pt') trainer.save_checkpoint(initial_state_checkpoint, {'epoch': 0}) batches_per_epoch = args.mdl_batches_per_epoch batch_size = args.mdl_batch_size block_size = args.mdl_block_size epoch_itr = trainer.get_train_iterator(epoch=0, load_dataset=True) examples = list(range(len(epoch_itr.dataset))) train_examples = examples[:args.mdl_train_examples] test_examples = examples[args.mdl_train_examples:] random.shuffle(test_examples) blocks = [train_examples] blocks += [ test_examples[i:i + block_size] for i in range(0, len(test_examples), block_size) ] allowed_examples = [] steps = len(blocks) block_cross_entropys = [] for step in range(steps): trainer.load_checkpoint(initial_state_checkpoint, reset_optimizer=True, reset_lr_scheduler=True) epoch_itr = trainer.get_train_iterator(epoch=step, load_dataset=False) allowed_examples += blocks[step] # if mdl-batch-size is set, we sample batches with replacement, # otherwise, each batch contains all allowed_examples if batch_size: batches = tuple([ random.choices(allowed_examples, k=batch_size) for _ in range(batches_per_epoch) ]) else: batches = tuple( [allowed_examples for _ in range(batches_per_epoch)]) epoch_itr.frozen_batches = batches train(args, trainer, task, epoch_itr) stashed_criterion = trainer.criterion train.criterion = CRITERION_REGISTRY['cross_entropy'](args, task) if step < steps - 1: stashed_criterion = trainer.criterion train.criterion = CRITERION_REGISTRY['cross_entropy'](args, task) next_block = (blocks[step + 1], ) next_block_cross_entropy = validate(args, trainer, task, epoch_itr, subsets=['train'], \ allowed_batches=next_block) train.criterion = stashed_criterion block_cross_entropys.append(next_block_cross_entropy) trainer.set_num_updates( 0 ) #reset the num_update as not systematically updated in load_checkpoint state_checkpoint = str(pathlib.Path(args.save_dir) / f'{step}.pt') trainer.save_checkpoint(state_checkpoint, {'epoch': step}) examples_seen = [len(b) for b in blocks] cross_entropy_sum = sum(n_examples * mean_cross_entropy for n_examples, mean_cross_entropy in zip( examples_seen[1:], block_cross_entropys)) stats = dict(online_cross_entropy=block_cross_entropys, description_length=cross_entropy_sum, examples_seen=examples_seen) print(json.dumps(stats)) state_checkpoint = str(pathlib.Path(args.save_dir) / 'last.pt') trainer.save_checkpoint(state_checkpoint, {'epoch': step})
def main(args): print(args) if not torch.cuda.is_available(): raise NotImplementedError('Training on CPU is not supported') torch.cuda.set_device(args.device_id) torch.manual_seed(args.seed) # Load dataset splits = ['train', 'valid'] if data.has_binary_files(args.data, splits): dataset = data.load_dataset(args.data, splits, args.source_lang, args.target_lang) else: dataset = data.load_raw_text_dataset(args.data, splits, args.source_lang, args.target_lang) if args.source_lang is None or args.target_lang is None: # record inferred languages in args, so that it's saved in checkpoints args.source_lang, args.target_lang = dataset.src, dataset.dst print('| [{}] dictionary: {} types'.format(dataset.src, len(dataset.src_dict))) print('| [{}] dictionary: {} types'.format(dataset.dst, len(dataset.dst_dict))) for split in splits: print('| {} {} {} examples'.format(args.data, split, len(dataset.splits[split]))) # Build model and criterion model = models.build_model(args, dataset.src_dict, dataset.dst_dict) if 0.0 < args.rank_scale < 1.0: patch_transformer(args, model) if args.wd2fd: no_decay, skiplist = [ 'fc1', 'fc2', 'embed_tokens', 'embed_positions', 'out_embed' ], [] else: no_decay, skiplist = [], [ 'fc1', 'fc2', 'embed_tokens', 'embed_positions', 'out_embed' ] else: no_decay, skiplist = [], [] spectral_init(args, model) criterion = criterions.build_criterion(args, dataset.src_dict, dataset.dst_dict) print('| model {}, criterion {}'.format(args.arch, criterion.__class__.__name__)) print('| num. model params: {}'.format( sum(p.data.numel() for p in model.parameters()))) # Build trainer no_decay, skiplist = [], [] if args.wd2fd_quekey: no_decay.extend(['_query.weight', '_key.weight']) else: skiplist.append('quekey') if args.wd2fd_outval: no_decay.extend(['_value.weight', 'output_perform.weight']) else: skiplist.append('outval') trainer = Trainer(args, model, criterion, skiplist=skiplist, no_decay=no_decay) print('| training on {} GPUs'.format(args.distributed_world_size)) print('| max tokens per GPU = {} and max sentences per GPU = {}'.format( args.max_tokens, args.max_sentences, )) # Load the latest checkpoint if one is available os.makedirs(args.save_dir, exist_ok=True) checkpoint_path = os.path.join(args.save_dir, args.restore_file) extra_state = trainer.load_checkpoint(checkpoint_path) if extra_state is not None: epoch = extra_state['epoch'] batch_offset = extra_state['batch_offset'] print('| loaded checkpoint {} (epoch {})'.format( checkpoint_path, epoch)) if batch_offset == 0: trainer.lr_step(epoch) epoch += 1 else: epoch, batch_offset = 1, 0 if args.distributed_rank <= 0: writer = SummaryWriter(args.save_dir) with open(os.path.join(args.save_dir, 'args.json'), 'w') as f: json.dump(vars(args), f, indent=4) else: writer = SummaryWriter( os.path.join(args.save_dir, str(args.distributed_rank))) # Train until the learning rate gets too small max_epoch = args.max_epoch or math.inf max_update = args.max_update or math.inf lr = trainer.get_lr() train_meter = StopwatchMeter() train_meter.start() while lr > args.min_lr and epoch <= max_epoch: if args.distributed_rank <= 0: writer.add_scalar('hyper/lr', lr, epoch) for form in ['QueKey', 'OutVal']: frobnorm, nucnorm, bound, nonorth = [], [], [], [] for module in model.modules(): if hasattr(module, form.lower()): U, VT = getattr(module, form.lower()).get_UVT() for u, vt in zip(U, VT): frobnorm.append(frobenius_norm(u, vt)) nucnorm.append( torch.norm(torch.matmul(u, vt), 'nuc')) bound.append( (u.pow(2).sum() + vt.pow(2).sum()) / 2.) nonorth.append(sum(non_orthogonality(u, vt)) / 2.) writer.add_scalar('FrobNorm/' + form, sum(frobnorm) / len(frobnorm), epoch) writer.add_scalar('NucNorm/' + form, sum(nucnorm) / len(nucnorm), epoch) writer.add_scalar('NucNorm/' + form + '-Bound', sum(bound) / len(bound), epoch) writer.add_scalar('NonOrth/' + form, sum(nonorth) / len(nonorth), epoch) frobnorm, nucnorm, bound, nonorth = [], [], [], [] for name, module in model.named_modules(): if not any( block in name for block in ['embed', '_query', '_key', '_value', 'output_perform']): if hasattr(module, 'frobgrad') and not hasattr(module, 'get_UVT'): U, VT = module.U.data, module.VT.data frobnorm.append(frobenius_norm(U, VT)) nucnorm.append(torch.norm(torch.matmul(U, VT), 'nuc')) nonorth.append(sum(non_orthogonality(U, VT)) / 2.) bound.append((U.pow(2).sum() + VT.pow(2).sum()) / 2.) elif hasattr(module, 'weight'): frobnorm.append(torch.norm(module.weight.data)) nucnorm.append(torch.norm(module.weight.data, 'nuc')) writer.add_scalar('FrobNorm/Linear', sum(frobnorm) / len(frobnorm), epoch) writer.add_scalar('NucNorm/Linear', sum(nucnorm) / len(nucnorm), epoch) if nonorth: writer.add_scalar('NucNorm/Linear-Bound', sum(bound) / len(bound), epoch) writer.add_scalar('NonOrth/Linear', sum(nonorth) / len(nonorth), epoch) # train for one epoch train(args, trainer, dataset, epoch, batch_offset) # evaluate on validate set if epoch % args.validate_interval == 0: for k, subset in enumerate(args.valid_subset.split(',')): val_loss = validate(args, trainer, dataset, subset, epoch) if k == 0: # only use first validation loss to update the learning schedule lr = trainer.lr_step(epoch, val_loss) # save checkpoint if not args.no_save: save_checkpoint(trainer, args, epoch, 0, val_loss) for k in ['loss', 'nll_loss']: writer.add_scalar('valid/' + k, trainer.meters['valid_' + k].avg, epoch) writer.add_scalar('train/' + k, trainer.meters['train_' + k].avg, epoch) else: lr = trainer.lr_step(epoch) epoch += 1 batch_offset = 0 if trainer.get_num_updates() >= max_update: break train_meter.stop() print('| done training in {:.1f} seconds'.format(train_meter.sum)) writer.flush() newpar = sum(p.numel() for p in model.parameters()) if 0.0 < args.rank_scale < 1.0: args.rank_scale = 1.0 origpar = sum(p.numel() for p in models.build_model( args, dataset.src_dict, dataset.dst_dict).parameters()) else: origpar = newpar if args.distributed_rank <= 0: with open(os.path.join(args.save_dir, 'results.json'), 'w') as f: json.dump( { 'final validation loss': trainer.meters['valid_nll_loss'].avg, 'original parameter count': origpar, 'compressed parameter count': newpar, 'compression ratio': newpar / origpar }, f, indent=4)