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)
