def __init__(self, profiler_type, logdir=Path('.')):
if profiler_type == 'None':
enabled = False
nvtx = False
elif profiler_type == 'CPU':
enabled = True
nvtx = False
elif profiler_type == 'NVTX':
enabled = True
nvtx = True
else:
assert False, f'Unknown profiler type {profiler_type}'
self._prof = None
self._enabled = enabled
self._nvtx = nvtx
self._logdir = Path(logdir)
self._logdir.mkdir(exist_ok=True, parents=True)
major, minor, _ = get_torch_version()
self._prof_kwargs = dict(enabled=self._enabled, profile_memory=True)
self._group_kwargs = dict()
if major > 1 or major == 1 and minor > 6:
self._prof_kwargs.update(dict(with_stack=True))
self._group_kwargs.update(dict(group_by_stack_n=15))
global is_pyprof_initialized
if not is_pyprof_initialized and self._enabled:
pyprof.init()
is_pyprof_initialized = True
def main():
args = parseArgs()
pyprof.init()
N = args.b
C = 3
H = d[args.m]['H']
W = d[args.m]['W']
opts = d[args.m]['opts']
classes = 1000
net = getattr(models, args.m)
net = net(**opts).cuda().half()
net.train()
x = torch.rand(N, C, H, W).cuda().half()
target = torch.empty(N, dtype=torch.long).random_(classes).cuda()
criterion = nn.CrossEntropyLoss().cuda()
if (args.o == "sgd"):
optimizer = torch.optim.SGD(net.parameters(), lr=0.01, momentum=0.9)
elif (args.o == "adam"):
optimizer = FusedAdam(net.parameters())
else:
assert False
#Warm up without profiler
for i in range(2):
output = net(x)
loss = criterion(output, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
with torch.autograd.profiler.emit_nvtx():
profiler.start()
output = net(x)
loss = criterion(output, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
profiler.stop()
# eval_punct_labels,
# eval_capit_labels,
# eval_subtokens_mask],
# user_iter_callback=lambda x, y: eval_iter_callback(x, y),
# user_epochs_done_callback=lambda x: eval_epochs_done_callback(x,
# punct_label_ids,
# capit_label_ids),
# eval_step=steps_per_epoch)
# Callback to store checkpoints
#ckpt_callback = nemo.core.CheckpointCallback(
# folder=nf.checkpoint_dir,
# step_freq=STEP_FREQ)
lr_policy = WarmupAnnealing(NUM_EPOCHS * steps_per_epoch,
warmup_ratio=LR_WARMUP_PROPORTION)
pyprof.init()
with torch.autograd.profiler.emit_nvtx():
nf.train(tensors_to_optimize=[task_loss],
callbacks=[callback_train],
lr_policy=lr_policy,
batches_per_step=BATCHES_PER_STEP,
optimizer=OPTIMIZER,
optimization_params={
"num_epochs": NUM_EPOCHS,
"lr": LEARNING_RATE
})
def main():
args = parse_args()
if args.affinity != 'disabled':
nproc_per_node = torch.cuda.device_count()
affinity = utils.gpu_affinity.set_affinity(args.local_rank,
nproc_per_node,
args.affinity)
print(f'{args.local_rank}: thread affinity: {affinity}')
# Initialize device and distributed backend
torch.cuda.set_device(args.local_rank)
l2_promote()
device = torch.device('cuda' if args.cuda else 'cpu')
utils.distributed.init_distributed(args.cuda)
args.work_dir = utils.exp_utils.build_work_dir_name(
args.work_dir,
args.dataset,
args.append_dataset,
args.append_time,
)
with utils.distributed.sync_workers() as rank:
if rank == 0:
create_exp_dir(args.work_dir,
scripts_to_save=['train.py', 'mem_transformer.py'],
debug=args.debug)
# Setup logging
if args.log_all_ranks:
log_file = f'train_log_rank_{utils.distributed.get_rank()}.log'
else:
log_file = args.txtlog_file
dllog_file = args.dllog_file
log_file = os.path.join(args.work_dir, log_file)
dllog_file = os.path.join(args.work_dir, dllog_file)
if args.debug:
log_file = os.devnull
dllog_file = os.devnull
utils.exp_utils.setup_logging(
log_all_ranks=args.log_all_ranks,
filename=log_file,
)
utils.exp_utils.setup_dllogger(enabled=True, filename=dllog_file)
if args.local_batch_size is not None:
world_size = utils.distributed.get_world_size()
args.batch_size = world_size * args.local_batch_size
logging.info(f'--local_batch_size was set, adjusting global batch size'
f' to {args.batch_size} (local_batch_size * world_size)')
if args.batch_size % args.batch_chunk != 0:
raise RuntimeError('Batch size needs to be divisible by '
'batch chunk')
if args.profile:
try:
pyprof.init(enable_function_stack=True)
except NameError:
warnings.warn('Called pyprof.init() but pyprof is not available')
logging.info(args)
dllogger.log(step='PARAMETER', data=vars(args))
logging.info(f'world size: {utils.distributed.get_world_size()}')
if not args.no_env:
log_env_info()
register_ignoring_timeout_handler()
# Set the random seed manually for reproducibility.
np.random.seed(args.seed)
torch.manual_seed(args.seed)
###########################################################################
# Load data
###########################################################################
corpus = get_lm_corpus(args.data, args.dataset, args.vocab)
ntokens = len(corpus.vocab)
vocab = corpus.vocab
args.n_token = ntokens
if args.mem_len == 0:
eval_mem_len = 0
else:
eval_mem_len = args.mem_len + args.tgt_len - args.eval_tgt_len
tr_iter = corpus.get_iterator('train',
args.batch_size,
args.tgt_len,
device=device,
ext_len=args.ext_len)
va_iter = corpus.get_iterator('valid',
args.eval_batch_size,
args.eval_tgt_len,
device=device,
mem_len=eval_mem_len,
ext_len=args.ext_len)
te_iter = corpus.get_iterator('test',
args.eval_batch_size,
args.eval_tgt_len,
device=device,
mem_len=eval_mem_len,
ext_len=args.ext_len)
# adaptive softmax / embedding
cutoffs, tie_projs = [], [False]
if args.adaptive:
assert args.dataset in ['wt103', 'lm1b']
if args.dataset == 'wt103':
cutoffs = [19997, 39997, 199997]
tie_projs += [True] * len(cutoffs)
elif args.dataset == 'lm1b':
cutoffs = [59997, 99997, 639997]
tie_projs += [False] * len(cutoffs)
###########################################################################
# Build the model
###########################################################################
model_config = {
'n_token': ntokens,
'n_layer': args.n_layer,
'n_head': args.n_head,
'd_model': args.d_model,
'd_head': args.d_head,
'd_inner': args.d_inner,
'dropout': args.dropout,
'dropatt': args.dropatt,
'dtype': None,
'tie_weight': args.tied,
'd_embed': args.d_embed,
'div_val': args.div_val,
'tie_projs': tie_projs,
'pre_lnorm': args.pre_lnorm,
'tgt_len': args.tgt_len,
'ext_len': args.ext_len,
'mem_len': args.mem_len,
'cutoffs': cutoffs,
'same_length': args.same_length,
'attn_type': args.attn_type,
'clamp_len': args.clamp_len,
'sample_softmax': args.sample_softmax,
}
model = MemTransformerLM(**model_config)
model.apply(functools.partial(weights_init, args=args))
# ensure embedding init is not overridden by out_layer in case of weight sharing
model.word_emb.apply(functools.partial(weights_init, args=args))
args.n_all_param = sum([p.nelement() for p in model.parameters()])
args.n_nonemb_param = sum(
[p.nelement() for p in model.layers.parameters()])
# optimizer
if args.optim.lower() == 'sgd':
if args.sample_softmax > 0:
dense_params, sparse_params = [], []
for param in model.parameters():
if param.size() == model.word_emb.weight.size():
sparse_params.append(param)
else:
dense_params.append(param)
optimizer_sparse = optim.SGD(sparse_params, lr=args.lr * 2)
optimizer = optim.SGD(dense_params, lr=args.lr, momentum=args.mom)
else:
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.mom)
optimizer_sparse = None
elif args.optim.lower() == 'adam':
if args.sample_softmax > 0:
dense_params, sparse_params = [], []
for param in model.parameters():
if param.size() == model.word_emb.weight.size():
sparse_params.append(param)
else:
dense_params.append(param)
optimizer_sparse = optim.SparseAdam(sparse_params, lr=args.lr)
optimizer = optim.Adam(dense_params,
lr=args.lr,
weight_decay=args.weight_decay)
else:
optimizer = optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
optimizer_sparse = None
elif args.optim.lower() == 'adagrad':
optimizer = optim.Adagrad(model.parameters(), lr=args.lr)
optimizer_sparse = None
elif args.optim.lower() == 'lamb':
optimizer = lamb.Lamb(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
optimizer_sparse = None
elif args.optim.lower() == 'jitlamb':
optimizer = lamb.JITLamb(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
optimizer_sparse = None
model = model.to(device)
scaler = None
if args.fp16:
if args.amp == 'pytorch':
scaler = torch.cuda.amp.GradScaler()
elif args.amp == 'apex':
model, optimizer = amp.initialize(
model,
optimizer,
opt_level=args.apex_amp_opt_level,
)
if args.multi_gpu == 'ddp' and torch.distributed.is_initialized():
para_model = DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,
broadcast_buffers=False,
find_unused_parameters=True,
)
elif args.multi_gpu == 'dp':
if args.gpu0_bsz >= 0:
para_model = BalancedDataParallel(args.gpu0_bsz //
args.batch_chunk,
model,
dim=1).to(device)
else:
para_model = nn.DataParallel(model, dim=1).to(device)
else:
para_model = model
# scheduler
if args.scheduler == 'cosine':
if args.max_step_scheduler:
max_step = args.max_step_scheduler
else:
max_step = args.max_step
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer,
max_step -
args.warmup_step,
eta_min=args.eta_min)
if args.sample_softmax > 0 and optimizer_sparse is not None:
scheduler_sparse = optim.lr_scheduler.CosineAnnealingLR(
optimizer_sparse,
max_step - args.warmup_step,
eta_min=args.eta_min)
else:
scheduler_sparse = None
elif args.scheduler == 'inv_sqrt':
# originally used for Transformer (in Attention is all you need)
def lr_lambda(step):
# return a multiplier instead of a learning rate
if step == 0 and args.warmup_step == 0:
return 1.
else:
return 1. / (step ** 0.5) if step > args.warmup_step \
else step / (args.warmup_step ** 1.5)
scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda)
if args.sample_softmax > 0 and optimizer_sparse is not None:
scheduler_sparse = optim.lr_scheduler.LambdaLR(optimizer_sparse,
lr_lambda=lr_lambda)
else:
scheduler_sparse = None
elif args.scheduler == 'dev_perf':
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer,
factor=args.decay_rate,
patience=args.patience,
min_lr=args.lr_min,
)
if args.sample_softmax > 0 and optimizer_sparse is not None:
scheduler_sparse = optim.lr_scheduler.ReduceLROnPlateau(
optimizer_sparse,
factor=args.decay_rate,
patience=args.patience,
min_lr=args.lr_min,
)
else:
scheduler_sparse = None
elif args.scheduler == 'constant':
pass
logging.info('=' * 100)
for k, v in args.__dict__.items():
logging.info(' - {} : {}'.format(k, v))
logging.info('=' * 100)
logging.info('#params = {}'.format(args.n_all_param))
logging.info('#non emb params = {}'.format(args.n_nonemb_param))
train_step = 0
start_epoch = 1
last_batch = 0
last_iter = 0
best_val_loss = None
if args.restart:
try:
checkpoint = load_checkpoint(args.restart)
model.load_state_dict(checkpoint['model_state'])
optimizer.load_state_dict(checkpoint['optimizer_state'])
scheduler.load_state_dict(checkpoint['scheduler_state'])
if args.fp16:
if args.amp == 'pytorch':
scaler.load_state_dict(checkpoint['amp_state'])
elif args.amp == 'apex':
amp.load_state_dict(checkpoint['amp_state'])
train_step = checkpoint['train_step']
start_epoch = checkpoint['epoch']
last_batch = checkpoint['batch']
last_iter = checkpoint['last_iter']
best_val_loss = checkpoint['best_val_loss']
if train_step >= args.max_step:
logging.info(
f'Loaded checkpoint after {train_step} steps, but '
f'this run was scheduled for a total of '
f'{args.max_step} steps, exiting')
sys.exit(1)
model.apply(functools.partial(update_dropout, args=args))
model.apply(functools.partial(update_dropatt, args=args))
except FileNotFoundError:
logging.info(f'Could not load checkpoint from {args.restart}, '
f'starting training from random init')
meters = {}
warmup = args.mem_len // args.tgt_len + 2
meters['train_throughput'] = AverageMeter(warmup=warmup)
###########################################################################
# Train
###########################################################################
# Loop over epochs.
# At any point you can hit Ctrl + C to break out of training early.
start_time = time.time()
with torch.autograd.profiler.emit_nvtx(enabled=args.profile):
with TimeoutHandler() as timeout_handler:
try:
for epoch in itertools.count(start=start_epoch):
if args.roll:
tr_iter.roll(seed=args.seed + epoch)
train_step, best_val_loss = train(
tr_iter, va_iter, model, para_model, model_config,
optimizer, optimizer_sparse, scheduler,
scheduler_sparse, scaler, vocab, epoch, last_batch,
last_iter, train_step, best_val_loss, meters,
timeout_handler, device, args)
last_batch = 0
last_iter = 0
if train_step == args.max_step:
logging.info('-' * 100)
logging.info('End of training')
break
except KeyboardInterrupt:
logging.info('-' * 100)
logging.info('Exiting from training early')
elapsed = time.time() - start_time
###########################################################################
# Test
###########################################################################
summary = {}
test_path = os.path.join(args.work_dir, 'checkpoint_best.pt')
if not args.debug and not args.no_eval and os.path.exists(test_path):
# Load the best saved model.
checkpoint = load_checkpoint(test_path)
model.load_state_dict(checkpoint['model_state'])
# Run on test data.
test_start_time = time.time()
with torch.autograd.profiler.emit_nvtx(enabled=args.profile):
test_loss = evaluate(te_iter, model, args)
test_loss = utils.distributed.all_reduce_item(test_loss, 'mean')
test_elapsed = time.time() - test_start_time
logging.info('=' * 100)
if args.dataset in ['enwik8', 'text8']:
logging.info(
'| End of training | test time: {:5.2f}s | test loss {:5.2f} | test bpc {:9.5f}'
.format(test_elapsed, test_loss, test_loss / math.log(2)))
else:
logging.info(
'| End of training | test time: {:5.2f}s | test loss {:5.2f} | test ppl {:9.3f}'
.format(test_elapsed, test_loss, math.exp(test_loss)))
logging.info('=' * 100)
summary.update({
'test_elapsed': test_elapsed,
'test_loss': test_loss,
})
if args.dataset in ['enwik8', 'text8']:
summary['test_bits_per_character'] = test_loss / math.log(2)
else:
summary['test_perplexity'] = math.exp(test_loss)
logging.info(f'Training time: {(elapsed / 60):.2f} minutes')
logging.info(
f'Training throughput: {meters["train_throughput"].avg:.2f} tok/s')
if best_val_loss:
val_perplexity = math.exp(best_val_loss)
else:
val_perplexity = None
summary.update({
'train_throughput': meters['train_throughput'].avg,
'train_elapsed': elapsed / 60,
'valid_loss': best_val_loss,
'valid_perplexity': val_perplexity,
})
dllogger.log(step=tuple(), data=summary)
passed = benchmark(target_perplexity=args.target_perplexity,
test_perplexity=val_perplexity,
target_throughput=args.target_throughput,
test_throughput=meters['train_throughput'].avg)
if not passed:
sys.exit(1)
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.optim
import torch.multiprocessing as mp
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
import torch.cuda.profiler as profiler
import pyprof
pyprof.init(enable_function_stack=True)
model_names = sorted(name for name in models.__dict__
if name.islower() and not name.startswith("__")
and callable(models.__dict__[name]))
parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
parser.add_argument('data', metavar='DIR', help='path to dataset')
parser.add_argument('-a',
'--arch',
metavar='ARCH',
default='resnet18',
choices=model_names,
help='model architecture: ' + ' | '.join(model_names) +
' (default: resnet18)')
parser.add_argument('-j',
def main():
args = parse_args()
if args.affinity != 'disabled':
nproc_per_node = torch.cuda.device_count()
affinity = utils.gpu_affinity.set_affinity(args.local_rank,
nproc_per_node,
args.affinity)
print(f'{args.local_rank}: thread affinity: {affinity}')
if args.type == 'pytorch':
from mem_transformer import MemTransformerLM
else:
from inference.mem_transformer_jit import MemTransformerLM
torch.cuda.set_device(args.local_rank)
l2_promote()
device = torch.device('cuda' if args.cuda else 'cpu')
utils.distributed.init_distributed(args.cuda)
with utils.distributed.sync_workers() as rank:
if rank == 0:
create_exp_dir(args.work_dir, debug=args.debug)
# Setup logging
if args.log_all_ranks:
log_file = f'eval_log_rank_{utils.distributed.get_rank()}.log'
else:
log_file = f'eval_log.log'
dllog_file = args.dllog_file
log_file = os.path.join(args.work_dir, log_file)
dllog_file = os.path.join(args.work_dir, dllog_file)
if args.debug:
log_file = os.devnull
dllog_file = os.devnull
utils.exp_utils.setup_logging(
log_all_ranks=args.log_all_ranks,
filename=log_file,
filemode='a',
)
utils.exp_utils.setup_dllogger(enabled=True, filename=dllog_file)
if args.profile:
try:
pyprof.init(enable_function_stack=True)
except NameError:
warnings.warn('Called pyprof.init() but pyprof is not available')
logging.info(args)
dllogger.log(step='PARAMETER', data=vars(args))
if not args.no_env:
log_env_info()
# Set the random seed manually for reproducibility.
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.model:
model_path = args.model
elif args.work_dir:
model_path = os.path.join(args.work_dir, 'checkpoint_best.pt')
else:
raise RuntimeError(
'Specify path to checkpoint using --model or --work_dir')
if not args.manual_config:
checkpoint = load_checkpoint(model_path)
vocab_type = checkpoint['args'].vocab
else:
checkpoint = None
vocab_type = args.manual_vocab
if args.manual:
vocab = checkpoint['vocab']
if hasattr(vocab, 'sym2idx') and not hasattr(vocab, 'unk_idx'):
vocab.unk_idx = vocab.sym2idx['<unk>']
text = " ".join(args.manual)
tokenized = tokenize_raw(text)
symbols = vocab.tokenize(tokenized, add_eos=True)
tensor = vocab.convert_to_tensor(symbols)
iter = data_utils.LMOrderedIterator(tensor,
bsz=args.batch_size,
bptt=args.tgt_len,
device=device,
ext_len=args.ext_len,
warmup=False)
else:
# Load dataset
corpus = get_lm_corpus(args.data, args.dataset, vocab_type)
if args.split == 'valid' or args.split == 'test':
iter = corpus.get_iterator(args.split,
args.batch_size,
args.tgt_len,
device=device,
mem_len=args.mem_len,
ext_len=args.ext_len)
else:
raise RuntimeError('Unknown split')
if args.fp16:
dtype = torch.float16
math_str = 'fp16'
else:
dtype = torch.float32
math_str = 'fp32'
if args.load_torchscript:
model = torch.jit.load(args.load_torchscript)
elif not args.manual_config:
checkpoint['model_config']['tgt_len'] = args.tgt_len
checkpoint['model_config']['ext_len'] = args.ext_len
checkpoint['model_config']['mem_len'] = args.mem_len
checkpoint['model_config']['clamp_len'] = args.clamp_len
checkpoint['model_config']['same_length'] = args.same_length
checkpoint['model_config']['dtype'] = dtype
model = MemTransformerLM(**checkpoint['model_config'])
if args.type == 'pytorch':
model.load_state_dict(checkpoint['model_state'])
elif args.type == 'torchscript':
model.load_state_dict(checkpoint['model_state'], strict=False)
elif args.manual_config:
args.manual_config['tgt_len'] = args.tgt_len
args.manual_config['ext_len'] = args.ext_len
args.manual_config['mem_len'] = args.mem_len
args.manual_config['clamp_len'] = args.clamp_len
args.manual_config['same_length'] = args.same_length
args.manual_config['dtype'] = dtype
model = MemTransformerLM(**args.manual_config)
model = model.eval()
model = model.to(device)
model = model.to(dtype)
if args.type == 'torchscript' and not args.manual_config:
state = checkpoint['model_state']
tie_projs = checkpoint['model_config']['tie_projs']
tie_weight = checkpoint['model_config']['tie_weight']
div_val = checkpoint['model_config']['div_val']
d_model = checkpoint['model_config']['d_model']
d_embed = checkpoint['model_config']['d_embed']
if div_val != 1 or d_model != d_embed:
for i in range(len(model.word_emb.emb_projs)):
model.word_emb.emb_projs[i] = state[
f'word_emb.emb_projs.{i}'].to(dtype)
for i in range(len(model.crit.out_projs)):
if div_val == 1:
src = 0
else:
src = i
if model.crit.out_projs[i] is not None:
if tie_projs[i]:
model.crit.out_projs[i] = state[
f'word_emb.emb_projs.{src}'].to(dtype)
else:
model.crit.out_projs[i] = state[f'crit.out_projs.{i}'].to(
dtype)
for i in range(len(model.crit.out_layers_biases)):
model.crit.out_layers_biases[i] = state[
f'crit.out_layers_biases.{i}'].to(dtype)
if tie_weight:
for i in range(len(model.crit.out_layers_weights)):
model.crit.out_layers_weights[i] = state[
f'word_emb.emb_layers.{i}.weight'].to(dtype)
else:
for i in range(len(model.crit.out_layers_weights)):
model.crit.out_layers_weights[i] = state[
f'crit.out_layers_weights.{i}'].to(dtype)
model = torch.jit.script(model)
if args.type != 'pytorch':
compile_model(model, device, args)
if args.type == 'torchscript' and args.save_torchscript:
torch.jit.save(model, args.save_torchscript)
logging.info(f'Evaluating with: math {math_str} type {args.type} '
f'bsz {args.batch_size} tgt_len {args.tgt_len} '
f'ext_len {args.ext_len} mem_len {args.mem_len} '
f'clamp_len {args.clamp_len}')
meters = {}
warmup = args.mem_len // args.tgt_len + 2
meters['eval_throughput'] = AverageMeter(warmup=warmup,
keep=args.save_data)
meters['eval_latency'] = AverageMeter(warmup=warmup, keep=args.save_data)
with torch.autograd.profiler.emit_nvtx(enabled=args.profile):
loss = evaluate(iter, model, meters, args.log_interval, args.max_size,
args.repeat)
perplexity = math.exp(loss)
log_str = format_log(loss, args.split, args)
summary = {
'eval_loss': loss,
'eval_ppl': perplexity,
}
logging.info('=' * 100)
logging.info(log_str)
logging.info('=' * 100)
if args.save_data:
latency_data = np.array(meters['eval_latency'].vals)
throughput_data = np.array(meters['eval_throughput'].vals)
precision = 'fp16' if args.fp16 else 'fp32'
data_fname = f'eval_data_{args.batch_size}_{precision}_{args.type}'
data_path = os.path.join(args.work_dir, data_fname)
data = {
'args': args,
'throughput': throughput_data,
'latency': latency_data,
}
with open(data_path, 'wb') as f:
pickle.dump(data, f)
logging.info(f'Throughput Avg: {throughput_data.mean():.2f} tok/s')
logging.info(f'Latency Avg: {1000.0 * latency_data.mean():.2f} ms')
for p in args.percentiles:
logging.info(
f'Latency {p}%: {1000.0 * np.percentile(latency_data, p):.2f} ms'
)
logging.info('=' * 100)
summary.update({
'eval_throughput': throughput_data.mean(),
'eval_avg_latency': 1000 * latency_data.mean(),
})
for p in args.percentiles:
summary[f'eval_{p}%_latency'] = 1000 * np.percentile(
latency_data, p)
dllogger.log(step=tuple(), data=summary)
passed = benchmark(
target_perplexity=args.target_perplexity,
test_perplexity=perplexity,
target_throughput=args.target_throughput,
test_throughput=meters['eval_throughput'].avg,
)
if not passed:
sys.exit(1)
def main():
"""
Launches translation (inference).
Inference is executed on a single GPU, implementation supports beam search
with length normalization and coverage penalty.
"""
args = parse_args()
if args.affinity != 'disabled':
nproc_per_node = torch.cuda.device_count()
affinity = gpu_affinity.set_affinity(args.local_rank, nproc_per_node,
args.affinity)
print(f'{args.local_rank}: thread affinity: {affinity}')
device = utils.set_device(args.cuda, args.local_rank)
utils.init_distributed(args.cuda)
args.rank = utils.get_rank()
os.makedirs(args.save_dir, exist_ok=True)
utils.setup_logging()
dllog_file = os.path.join(args.save_dir, args.dllog_file)
utils.setup_dllogger(enabled=True, filename=dllog_file)
if args.profile:
try:
pyprof.init(enable_function_stack=True)
except NameError:
warnings.warn('Called pyprof.init() but pyprof is not available')
if args.env:
utils.log_env_info()
logging.info(f'Run arguments: {args}')
dllogger.log(step='PARAMETER', data=vars(args))
if not args.cuda and torch.cuda.is_available():
warnings.warn('cuda is available but not enabled')
if not args.cudnn:
torch.backends.cudnn.enabled = False
# load checkpoint and deserialize to CPU (to save GPU memory)
if args.model:
checkpoint = torch.load(args.model, map_location={'cuda:0': 'cpu'})
# build GNMT model
tokenizer = Tokenizer()
tokenizer.set_state(checkpoint['tokenizer'])
model_config = checkpoint['model_config']
model_config['batch_first'] = args.batch_first
model_config['vocab_size'] = tokenizer.vocab_size
model = GNMT(**model_config)
model.load_state_dict(checkpoint['state_dict'])
elif args.synthetic:
model = GNMT(args.synthetic_vocab, batch_first=args.batch_first)
tokenizer = None
else:
raise RuntimeError(
'Specify model either with --synthetic or with --model flag')
# construct the dataset
if args.input:
data = RawTextDataset(
raw_datafile=args.input,
tokenizer=tokenizer,
sort=args.sort,
)
elif args.input_text:
data = RawTextDataset(
raw_data=args.input_text,
tokenizer=tokenizer,
sort=args.sort,
)
elif args.synthetic:
data = SyntheticDataset(args.synthetic_vocab, args.synthetic_len,
args.batch_size[0] * args.synthetic_batches)
latency_table = tables.LatencyTable(args.percentiles)
throughput_table = tables.ThroughputTable(args.percentiles)
accuracy_table = tables.AccuracyTable('BLEU')
dtype = {
'fp32': torch.FloatTensor,
'tf32': torch.FloatTensor,
'fp16': torch.HalfTensor
}
for (math, batch_size, beam_size) in product(args.math, args.batch_size,
args.beam_size):
logging.info(f'math: {math}, batch size: {batch_size}, '
f'beam size: {beam_size}')
model.type(dtype[math])
model = model.to(device)
model.eval()
# build the data loader
loader = data.get_loader(
batch_size=batch_size,
batch_first=args.batch_first,
pad=True,
repeat=args.repeat[batch_size],
num_workers=0,
)
# build the translator object
translator = Translator(
model=model,
tokenizer=tokenizer,
loader=loader,
beam_size=beam_size,
max_seq_len=args.max_seq_len,
len_norm_factor=args.len_norm_factor,
len_norm_const=args.len_norm_const,
cov_penalty_factor=args.cov_penalty_factor,
print_freq=args.print_freq,
)
# execute the inference
with torch.autograd.profiler.emit_nvtx(enabled=args.profile):
output, stats = translator.run(
calc_bleu=args.bleu,
eval_path=args.output,
summary=True,
warmup=args.warmup,
reference_path=args.reference,
)
# print translated outputs
if not args.synthetic and (not args.output and args.rank == 0):
logging.info(f'Translated output:')
for out in output:
print(out)
key = (batch_size, beam_size)
latency_table.add(key, {math: stats['runtimes']})
throughput_table.add(key, {math: stats['throughputs']})
accuracy_table.add(key, {math: stats['bleu']})
if args.tables:
accuracy_table.write('Inference accuracy', args.math)
if 'fp16' in args.math and 'fp32' in args.math:
relative = 'fp32'
elif 'fp16' in args.math and 'tf32' in args.math:
relative = 'tf32'
else:
relative = None
if 'fp32' in args.math:
throughput_table.write('Inference throughput', 'fp32')
if 'tf32' in args.math:
throughput_table.write('Inference throughput', 'tf32')
if 'fp16' in args.math:
throughput_table.write('Inference throughput',
'fp16',
relative=relative)
if 'fp32' in args.math:
latency_table.write('Inference latency', 'fp32')
if 'tf32' in args.math:
latency_table.write('Inference latency', 'tf32')
if 'fp16' in args.math:
latency_table.write('Inference latency',
'fp16',
relative=relative,
reverse_speedup=True)
avg_throughput = np.array(stats['throughputs']).mean()
avg_latency = np.array(stats['runtimes']).mean()
summary = {
'eval_throughput': avg_throughput,
'eval_bleu': stats['bleu'],
'eval_avg_latency': avg_latency,
}
for p in args.percentiles:
summary[f'eval_{p}%_latency'] = np.percentile(stats['runtimes'], p)
dllogger.log(step=tuple(), data=summary)
passed = utils.benchmark(stats['bleu'], args.target_bleu,
stats['tokens_per_sec'], args.target_perf)
return passed
def main():
# Parse command line arguments
args = parse_args()
# DLProf - Init PyProf
if args.dlprof:
pyprof.init(enable_function_stack=True)
# Set num epochs to 1 if DLProf is enabled
args.epochs = 1
# Create experiment
summary = _make_experiment(args)
# Create datasets
train_data = KittiObjectDataset(args.root, 'train', args.grid_size,
args.grid_res, args.yoffset)
val_data = KittiObjectDataset(args.root, 'val', args.grid_size,
args.grid_res, args.yoffset)
# Apply data augmentation
# train_data = oft.AugmentedObjectDataset(
# train_data, args.train_image_size, args.train_grid_size,
# jitter=args.grid_jitter)
# Create dataloaders
train_loader = DataLoader(train_data,
args.batch_size,
shuffle=True,
num_workers=args.workers,
collate_fn=oft.utils.collate)
val_loader = DataLoader(val_data,
args.batch_size,
shuffle=False,
num_workers=args.workers,
collate_fn=oft.utils.collate)
# Build model
model = OftNet(num_classes=1,
frontend=args.frontend,
topdown_layers=args.topdown,
grid_res=args.grid_res,
grid_height=args.grid_height)
if len(args.gpu) > 0:
torch.cuda.set_device(args.gpu[0])
model = nn.DataParallel(model, args.gpu).cuda()
# Create encoder
encoder = ObjectEncoder()
# Setup optimizer
optimizer = optim.SGD(model.parameters(), args.lr, args.momentum,
args.weight_decay)
scheduler = optim.lr_scheduler.ExponentialLR(optimizer, args.lr_decay)
# Creates a GradScaler once at the beginning of training for AMP. Created even if not being used.
scaler = GradScaler()
for epoch in range(1, args.epochs + 1):
print('\n=== Beginning epoch {} of {} ==='.format(epoch, args.epochs))
# Update and log learning rate
scheduler.step(epoch - 1)
summary.add_scalar('lr', optimizer.param_groups[0]['lr'], epoch)
# Train model
if args.dlprof:
with torch.autograd.profiler.emit_nvtx():
train(args, train_loader, model, encoder, optimizer, summary,
epoch, scaler)
else:
train(args, train_loader, model, encoder, optimizer, summary,
epoch, scaler)
# Run validation every N epochs
if epoch % args.val_interval == 0:
if args.dlprof:
with torch.autograd.profiler.emit_nvtx():
validate(args, val_loader, model, encoder, summary, epoch)
else:
validate(args, val_loader, model, encoder, summary, epoch)
# Save model checkpoint
save_checkpoint(args, epoch, model, optimizer, scheduler)
import numpy as np
import torch
import time
from flowtron import Flowtron
from torch.utils.data import DataLoader
from data import Data
from train import update_params
sys.path.insert(0, "tacotron2")
sys.path.insert(0, "tacotron2/waveglow")
from glow import WaveGlow
from scipy.io.wavfile import write
pyprof.init() ########### prof.
def infer(flowtron_path, waveglow_path, text, speaker_id, n_frames, sigma,
seed):
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
# load waveglow
waveglow = torch.load(waveglow_path)['model'].cuda().eval()
waveglow.cuda().half()
for k in waveglow.convinv:
k.float()
waveglow.eval()
# load flowtron
def main():
args = parse_args()
assert (torch.cuda.is_available())
assert args.prediction_frequency % args.log_frequency == 0
torch.backends.cudnn.benchmark = args.cudnn_benchmark
# set up distributed training
multi_gpu = int(os.environ.get('WORLD_SIZE', 1)) > 1
if multi_gpu:
torch.cuda.set_device(args.local_rank)
dist.init_process_group(backend='nccl', init_method='env://')
world_size = dist.get_world_size()
print_once(f'Distributed training with {world_size} GPUs\n')
else:
world_size = 1
torch.manual_seed(args.seed + args.local_rank)
np.random.seed(args.seed + args.local_rank)
random.seed(args.seed + args.local_rank)
init_log(args)
cfg = config.load(args.model_config)
config.apply_duration_flags(cfg, args.max_duration,
args.pad_to_max_duration)
symbols = helpers.add_ctc_blank(cfg['labels'])
assert args.grad_accumulation_steps >= 1
assert args.batch_size % args.grad_accumulation_steps == 0
batch_size = args.batch_size // args.grad_accumulation_steps
print_once('Setting up datasets...')
train_dataset_kw, train_features_kw = config.input(cfg, 'train')
val_dataset_kw, val_features_kw = config.input(cfg, 'val')
use_dali = args.dali_device in ('cpu', 'gpu')
if use_dali:
assert train_dataset_kw['ignore_offline_speed_perturbation'], \
"DALI doesn't support offline speed perturbation"
# pad_to_max_duration is not supported by DALI - have simple padders
if train_features_kw['pad_to_max_duration']:
train_feat_proc = BaseFeatures(
pad_align=train_features_kw['pad_align'],
pad_to_max_duration=True,
max_duration=train_features_kw['max_duration'],
sample_rate=train_features_kw['sample_rate'],
window_size=train_features_kw['window_size'],
window_stride=train_features_kw['window_stride'])
train_features_kw['pad_to_max_duration'] = False
else:
train_feat_proc = None
if val_features_kw['pad_to_max_duration']:
val_feat_proc = BaseFeatures(
pad_align=val_features_kw['pad_align'],
pad_to_max_duration=True,
max_duration=val_features_kw['max_duration'],
sample_rate=val_features_kw['sample_rate'],
window_size=val_features_kw['window_size'],
window_stride=val_features_kw['window_stride'])
val_features_kw['pad_to_max_duration'] = False
else:
val_feat_proc = None
train_loader = DaliDataLoader(
gpu_id=args.local_rank,
dataset_path=args.dataset_dir,
config_data=train_dataset_kw,
config_features=train_features_kw,
json_names=args.train_manifests,
batch_size=batch_size,
grad_accumulation_steps=args.grad_accumulation_steps,
pipeline_type="train",
device_type=args.dali_device,
symbols=symbols)
val_loader = DaliDataLoader(gpu_id=args.local_rank,
dataset_path=args.dataset_dir,
config_data=val_dataset_kw,
config_features=val_features_kw,
json_names=args.val_manifests,
batch_size=batch_size,
pipeline_type="val",
device_type=args.dali_device,
symbols=symbols)
else:
train_dataset_kw, train_features_kw = config.input(cfg, 'train')
train_dataset = AudioDataset(args.dataset_dir, args.train_manifests,
symbols, **train_dataset_kw)
train_loader = get_data_loader(train_dataset,
batch_size,
multi_gpu=multi_gpu,
shuffle=True,
num_workers=4)
train_feat_proc = FilterbankFeatures(**train_features_kw)
val_dataset_kw, val_features_kw = config.input(cfg, 'val')
val_dataset = AudioDataset(args.dataset_dir, args.val_manifests,
symbols, **val_dataset_kw)
val_loader = get_data_loader(val_dataset,
batch_size,
multi_gpu=multi_gpu,
shuffle=False,
num_workers=4,
drop_last=False)
val_feat_proc = FilterbankFeatures(**val_features_kw)
dur = train_dataset.duration / 3600
dur_f = train_dataset.duration_filtered / 3600
nsampl = len(train_dataset)
print_once(f'Training samples: {nsampl} ({dur:.1f}h, '
f'filtered {dur_f:.1f}h)')
if train_feat_proc is not None:
train_feat_proc.cuda()
if val_feat_proc is not None:
val_feat_proc.cuda()
steps_per_epoch = len(train_loader) // args.grad_accumulation_steps
# set up the model
model = Jasper(encoder_kw=config.encoder(cfg),
decoder_kw=config.decoder(cfg, n_classes=len(symbols)))
model.cuda()
ctc_loss = CTCLossNM(n_classes=len(symbols))
greedy_decoder = GreedyCTCDecoder()
print_once(f'Model size: {num_weights(model) / 10**6:.1f}M params\n')
# optimization
kw = {'lr': args.lr, 'weight_decay': args.weight_decay}
if args.optimizer == "novograd":
optimizer = Novograd(model.parameters(), **kw)
elif args.optimizer == "adamw":
optimizer = AdamW(model.parameters(), **kw)
else:
raise ValueError(f'Invalid optimizer "{args.optimizer}"')
adjust_lr = lambda step, epoch, optimizer: lr_policy(
step,
epoch,
args.lr,
optimizer,
steps_per_epoch=steps_per_epoch,
warmup_epochs=args.warmup_epochs,
hold_epochs=args.hold_epochs,
num_epochs=args.epochs,
policy=args.lr_policy,
min_lr=args.min_lr,
exp_gamma=args.lr_exp_gamma)
if args.amp:
model, optimizer = amp.initialize(min_loss_scale=1.0,
models=model,
optimizers=optimizer,
opt_level='O1',
max_loss_scale=512.0)
if args.ema > 0:
ema_model = copy.deepcopy(model)
else:
ema_model = None
if multi_gpu:
model = DistributedDataParallel(model)
if args.pyprof:
pyprof.init(enable_function_stack=True)
# load checkpoint
meta = {'best_wer': 10**6, 'start_epoch': 0}
checkpointer = Checkpointer(args.output_dir, 'Jasper',
args.keep_milestones, args.amp)
if args.resume:
args.ckpt = checkpointer.last_checkpoint() or args.ckpt
if args.ckpt is not None:
checkpointer.load(args.ckpt, model, ema_model, optimizer, meta)
start_epoch = meta['start_epoch']
best_wer = meta['best_wer']
epoch = 1
step = start_epoch * steps_per_epoch + 1
if args.pyprof:
torch.autograd.profiler.emit_nvtx().__enter__()
profiler.start()
# training loop
model.train()
# pre-allocate
if args.pre_allocate_range is not None:
n_feats = train_features_kw['n_filt']
pad_align = train_features_kw['pad_align']
a, b = args.pre_allocate_range
for n_frames in range(a, b + pad_align, pad_align):
print_once(
f'Pre-allocation ({batch_size}x{n_feats}x{n_frames})...')
feat = torch.randn(batch_size, n_feats, n_frames, device='cuda')
feat_lens = torch.ones(batch_size, device='cuda').fill_(n_frames)
txt = torch.randint(high=len(symbols) - 1,
size=(batch_size, 100),
device='cuda')
txt_lens = torch.ones(batch_size, device='cuda').fill_(100)
log_probs, enc_lens = model(feat, feat_lens)
del feat
loss = ctc_loss(log_probs, txt, enc_lens, txt_lens)
loss.backward()
model.zero_grad()
for epoch in range(start_epoch + 1, args.epochs + 1):
if multi_gpu and not use_dali:
train_loader.sampler.set_epoch(epoch)
epoch_utts = 0
accumulated_batches = 0
epoch_start_time = time.time()
for batch in train_loader:
if accumulated_batches == 0:
adjust_lr(step, epoch, optimizer)
optimizer.zero_grad()
step_loss = 0
step_utts = 0
step_start_time = time.time()
if use_dali:
# with DALI, the data is already on GPU
feat, feat_lens, txt, txt_lens = batch
if train_feat_proc is not None:
feat, feat_lens = train_feat_proc(feat, feat_lens,
args.amp)
else:
batch = [t.cuda(non_blocking=True) for t in batch]
audio, audio_lens, txt, txt_lens = batch
feat, feat_lens = train_feat_proc(audio, audio_lens, args.amp)
log_probs, enc_lens = model(feat, feat_lens)
loss = ctc_loss(log_probs, txt, enc_lens, txt_lens)
loss /= args.grad_accumulation_steps
if torch.isnan(loss).any():
print_once(f'WARNING: loss is NaN; skipping update')
else:
if multi_gpu:
step_loss += reduce_tensor(loss.data, world_size).item()
else:
step_loss += loss.item()
if args.amp:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
step_utts += batch[0].size(0) * world_size
epoch_utts += batch[0].size(0) * world_size
accumulated_batches += 1
if accumulated_batches % args.grad_accumulation_steps == 0:
optimizer.step()
apply_ema(model, ema_model, args.ema)
if step % args.log_frequency == 0:
preds = greedy_decoder(log_probs)
wer, pred_utt, ref = greedy_wer(preds, txt, txt_lens,
symbols)
if step % args.prediction_frequency == 0:
print_once(f' Decoded: {pred_utt[:90]}')
print_once(f' Reference: {ref[:90]}')
step_time = time.time() - step_start_time
log(
(epoch, step % steps_per_epoch
or steps_per_epoch, steps_per_epoch), step, 'train', {
'loss': step_loss,
'wer': 100.0 * wer,
'throughput': step_utts / step_time,
'took': step_time,
'lrate': optimizer.param_groups[0]['lr']
})
step_start_time = time.time()
if step % args.eval_frequency == 0:
wer = evaluate(epoch, step, val_loader, val_feat_proc,
symbols, model, ema_model, ctc_loss,
greedy_decoder, args.amp, use_dali)
if wer < best_wer and epoch >= args.save_best_from:
checkpointer.save(model,
ema_model,
optimizer,
epoch,
step,
best_wer,
is_best=True)
best_wer = wer
step += 1
accumulated_batches = 0
# end of step
# DALI iterator need to be exhausted;
# if not using DALI, simulate drop_last=True with grad accumulation
if not use_dali and step > steps_per_epoch * epoch:
break
epoch_time = time.time() - epoch_start_time
log((epoch, ), None, 'train_avg', {
'throughput': epoch_utts / epoch_time,
'took': epoch_time
})
if epoch % args.save_frequency == 0 or epoch in args.keep_milestones:
checkpointer.save(model, ema_model, optimizer, epoch, step,
best_wer)
if 0 < args.epochs_this_job <= epoch - start_epoch:
print_once(f'Finished after {args.epochs_this_job} epochs.')
break
# end of epoch
if args.pyprof:
profiler.stop()
torch.autograd.profiler.emit_nvtx().__exit__(None, None, None)
log((), None, 'train_avg', {'throughput': epoch_utts / epoch_time})
if epoch == args.epochs:
evaluate(epoch, step, val_loader, val_feat_proc, symbols, model,
ema_model, ctc_loss, greedy_decoder, args.amp, use_dali)
checkpointer.save(model, ema_model, optimizer, epoch, step, best_wer)
flush_log()
def main(argv):
validate_flags()
torch.manual_seed(FLAGS.seed)
utils.init_logging(log_path=FLAGS.log_path)
dllogger.log(data=FLAGS.flag_values_dict(), step='PARAMETER')
data_loader_train, data_loader_test = get_data_loaders(FLAGS)
scaled_lr = FLAGS.lr / FLAGS.loss_scale if FLAGS.amp else FLAGS.lr
model = create_model()
optimizer = torch.optim.SGD(model.parameters(), lr=scaled_lr)
if FLAGS.mode == 'prof-train':
pyprof.init(enable_function_stack=True)
if FLAGS.amp and (FLAGS.mode == 'train' or FLAGS.mode == 'prof-train'):
(model.top_model, model.bottom_model.mlp), optimizer = amp.initialize(
[model.top_model, model.bottom_model.mlp],
optimizer,
opt_level="O2",
loss_scale=1)
elif FLAGS.amp:
model = model.half()
loss_fn = torch.nn.BCEWithLogitsLoss(reduction="mean")
if FLAGS.mode == 'test':
loss, auc, test_step_time = evaluate(model, loss_fn, data_loader_test)
avg_test_throughput = FLAGS.batch_size / test_step_time
results = {
'auc': auc,
'avg_inference_latency': test_step_time,
'average_test_throughput': avg_test_throughput
}
dllogger.log(data=results, step=tuple())
print(f"Finished testing. Test Loss {loss:.4f}, auc {auc:.4f}")
return
if FLAGS.mode == 'inference_benchmark':
results = {}
if FLAGS.amp:
# can use pure FP16 for inference
model = model.half()
for batch_size in FLAGS.inference_benchmark_batch_sizes:
batch_size = int(batch_size)
FLAGS.test_batch_size = batch_size
_, benchmark_data_loader = get_data_loaders(FLAGS)
latencies = inference_benchmark(
model=model,
data_loader=benchmark_data_loader,
num_batches=FLAGS.inference_benchmark_steps)
print("All inference latencies: {}".format(latencies))
mean_latency = np.mean(latencies)
mean_inference_throughput = batch_size / mean_latency
subresult = {
f'mean_inference_latency_batch_{batch_size}':
mean_latency,
f'mean_inference_throughput_batch_{batch_size}':
mean_inference_throughput
}
results.update(subresult)
dllogger.log(data=results, step=tuple())
print(f"Finished inference benchmark.")
return
if FLAGS.mode == 'train':
train(model, loss_fn, optimizer, data_loader_train, data_loader_test,
scaled_lr)
if FLAGS.mode == 'prof-train':
with torch.autograd.profiler.emit_nvtx():
train(model, loss_fn, optimizer, data_loader_train,
data_loader_test, scaled_lr)
def main(args):
if args.profile:
pyprof.init(enable_function_stack=True)
print(args)
setup_logger(args)
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
if args.distributed_world_size > 1:
assert torch.distributed.is_initialized()
torch.distributed.broadcast(torch.tensor([1], device="cuda"), 0)
torch.cuda.synchronize()
pValue = ctypes.cast((ctypes.c_int * 1)(), ctypes.POINTER(ctypes.c_int))
ctypes.CDLL('libcudart.so').cudaDeviceSetLimit(ctypes.c_int(0x05),
ctypes.c_int(128))
ctypes.CDLL('libcudart.so').cudaDeviceGetLimit(pValue, ctypes.c_int(0x05))
torch.manual_seed(args.seed)
src_dict, tgt_dict = data_utils.load_dictionaries(args)
add_extra_items_to_checkpoint({'src_dict': src_dict, 'tgt_dict': tgt_dict})
datasets = load_dataset_splits(args, ['train', 'valid', 'test'], src_dict,
tgt_dict)
model = build_model(args)
print('| num. model params: {}'.format(
sum(p.numel() for p in model.parameters())))
# Build trainer
if torch.cuda.get_device_capability(0)[0] >= 7 and not args.amp:
print('| NOTICE: your device may support faster training with --amp')
trainer = DDPTrainer(args, model)
print('| model {}, criterion {}'.format(
args.arch, trainer.criterion.__class__.__name__))
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,
))
epoch_itr = data.EpochBatchIterator(
dataset=datasets[args.train_subset],
max_tokens=args.max_tokens,
max_sentences=args.max_sentences_valid,
max_positions=args.max_positions,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
)
# Load the latest checkpoint if one is available
load_checkpoint(args, trainer, epoch_itr)
# Send a dummy batch to warm the caching allocator
dummy_batch = data_utils.get_dummy_batch(args.max_tokens, src_dict,
tgt_dict)
trainer.dummy_train_step(dummy_batch)
# Sanity check
if args.do_sanity_check:
print('Performing sanity check...')
sanity_score = score(args, trainer, datasets['test'], src_dict,
tgt_dict, 'test.raw.de')
DLLogger.log(step='SANITY_CHECK',
data={'sanity_check_score': sanity_score},
verbosity=1)
# Train until the learning rate gets too small or model reaches target score
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
tgt_bleu = args.target_bleu or math.inf
current_bleu = 0.0
best_bleu = -1.0
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
run_summary = {
'loss': float('inf'),
'val_loss': float('inf'),
'speed': 0,
'accuracy': 0
}
while lr >= args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update and current_bleu < tgt_bleu:
DLLogger.log(step=trainer.get_num_updates() + 1,
data={'epoch': epoch_itr.epoch},
verbosity=0)
# train for one epoch
with torch.autograd.profiler.emit_nvtx(enabled=args.profile):
train(args, trainer, epoch_itr)
DLLogger.log(step=trainer.get_num_updates(),
data={'walltime': train_meter.sum},
verbosity=1)
DLLogger.log(step=trainer.get_num_updates(),
data={'avg_epoch_loss': trainer.avg_loss_meter.avg},
verbosity=1)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, datasets, valid_subsets)
valid_bleu = score(args, trainer, datasets[valid_subsets[0]],
src_dict, tgt_dict, 'valid.raw.de')
DLLogger.log(step=trainer.get_num_updates(),
data={
'val_loss': valid_losses[0],
'val_bleu': valid_bleu
},
verbosity=1)
# Eval BLEU score
if args.online_eval or (tgt_bleu is not math.inf):
current_bleu = score(args, trainer, datasets[args.gen_subset],
src_dict, tgt_dict, 'test.raw.de')
DLLogger.log(step=trainer.get_num_updates(),
data={'test_bleu': current_bleu},
verbosity=1)
best_bleu = max(best_bleu, current_bleu)
run_summary['val_loss'] = min(run_summary['val_loss'], valid_losses[0])
run_summary['accuracy'] = best_bleu if best_bleu >= 0 else valid_bleu
run_summary['loss'] = valid_losses[0]
run_summary['speed'] = trainer.throughput_meter.u_avg
# Only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
train_meter.stop()
run_summary['walltime'] = train_meter.sum
DLLogger.log(step=(), data=run_summary, verbosity=0)
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main():
"""
Launches data-parallel multi-gpu training.
"""
pyprof.init()
training_start = time.time()
args = parse_args()
device = utils.set_device(args.cuda, args.local_rank)
utils.init_distributed(args.cuda)
args.rank = utils.get_rank()
if not args.cudnn:
torch.backends.cudnn.enabled = False
# create directory for results
os.makedirs(args.save_dir, exist_ok=True)
# setup logging
log_filename = f'log_rank_{utils.get_rank()}.log'
utils.setup_logging(args.log_all_ranks,
os.path.join(args.save_dir, log_filename))
if args.env:
utils.log_env_info()
logging.info(f'Saving results to: {args.save_dir}')
logging.info(f'Run arguments: {args}')
args.train_iter_size = set_iter_size(args.train_iter_size,
args.train_global_batch_size,
args.train_batch_size)
worker_seeds, shuffling_seeds = utils.setup_seeds(args.seed, args.epochs,
device)
worker_seed = worker_seeds[args.rank]
logging.info(f'Worker {args.rank} is using worker seed: {worker_seed}')
torch.manual_seed(worker_seed)
# build tokenizer
pad_vocab = utils.pad_vocabulary(args.math)
tokenizer = Tokenizer(args.vocab, args.bpe_codes, args.lang, pad_vocab)
# build datasets
train_data = LazyParallelDataset(
src_fname=args.train_src,
tgt_fname=args.train_tgt,
tokenizer=tokenizer,
min_len=args.train_min_length,
max_len=args.train_max_length,
sort=False,
max_size=args.train_max_size,
)
val_data = ParallelDataset(
src_fname=args.val_src,
tgt_fname=args.val_tgt,
tokenizer=tokenizer,
min_len=args.val_min_length,
max_len=args.val_max_length,
sort=True,
)
test_data = TextDataset(
src_fname=args.test_src,
tokenizer=tokenizer,
min_len=args.test_min_length,
max_len=args.test_max_length,
sort=True,
)
vocab_size = tokenizer.vocab_size
# build GNMT model
model_config = {
'hidden_size': args.hidden_size,
'vocab_size': vocab_size,
'num_layers': args.num_layers,
'dropout': args.dropout,
'batch_first': False,
'share_embedding': args.share_embedding,
}
model = GNMT(**model_config).to(device)
logging.info(model)
batch_first = model.batch_first
# define loss function (criterion) and optimizer
criterion = build_criterion(vocab_size, config.PAD,
args.smoothing).to(device)
opt_config = {'optimizer': args.optimizer, 'lr': args.lr}
opt_config.update(literal_eval(args.optimizer_extra))
logging.info(f'Training optimizer config: {opt_config}')
scheduler_config = {
'warmup_steps': args.warmup_steps,
'remain_steps': args.remain_steps,
'decay_interval': args.decay_interval,
'decay_steps': args.decay_steps,
'decay_factor': args.decay_factor
}
logging.info(f'Training LR schedule config: {scheduler_config}')
num_parameters = sum([l.nelement() for l in model.parameters()])
logging.info(f'Number of parameters: {num_parameters}')
batching_opt = {
'shard_size': args.shard_size,
'num_buckets': args.num_buckets
}
# get data loaders
train_loader = train_data.get_loader(batch_size=args.train_batch_size,
seeds=shuffling_seeds,
batch_first=batch_first,
shuffle=True,
batching=args.batching,
batching_opt=batching_opt,
num_workers=args.train_loader_workers)
val_loader = val_data.get_loader(batch_size=args.val_batch_size,
batch_first=batch_first,
shuffle=False,
num_workers=args.val_loader_workers)
test_loader = test_data.get_loader(batch_size=args.test_batch_size,
batch_first=batch_first,
shuffle=False,
pad=True,
num_workers=args.test_loader_workers)
translator = Translator(
model=model,
tokenizer=tokenizer,
loader=test_loader,
beam_size=args.beam_size,
max_seq_len=args.test_max_length,
len_norm_factor=args.len_norm_factor,
len_norm_const=args.len_norm_const,
cov_penalty_factor=args.cov_penalty_factor,
print_freq=args.print_freq,
reference=args.test_tgt,
)
# create trainer
total_train_iters = len(train_loader) // args.train_iter_size * args.epochs
save_info = {
'model_config': model_config,
'config': args,
'tokenizer': tokenizer.get_state()
}
loss_scaling = {
'init_scale': args.init_scale,
'upscale_interval': args.upscale_interval
}
trainer_options = dict(
model=model,
criterion=criterion,
grad_clip=args.grad_clip,
iter_size=args.train_iter_size,
save_dir=args.save_dir,
save_freq=args.save_freq,
save_info=save_info,
opt_config=opt_config,
scheduler_config=scheduler_config,
train_iterations=total_train_iters,
keep_checkpoints=args.keep_checkpoints,
math=args.math,
loss_scaling=loss_scaling,
print_freq=args.print_freq,
intra_epoch_eval=args.intra_epoch_eval,
translator=translator,
prealloc_mode=args.prealloc_mode,
)
trainer = trainers.Seq2SeqTrainer(**trainer_options)
# optionally resume from a checkpoint
if args.resume:
checkpoint_file = args.resume
if os.path.isdir(checkpoint_file):
checkpoint_file = os.path.join(checkpoint_file, 'model_best.pth')
if os.path.isfile(checkpoint_file):
trainer.load(checkpoint_file)
else:
logging.error(f'No checkpoint found at {args.resume}')
# training loop
best_loss = float('inf')
training_perf = []
break_training = False
test_bleu = None
for epoch in range(args.start_epoch, args.epochs):
logging.info(f'Starting epoch {epoch}')
train_loader.sampler.set_epoch(epoch)
trainer.epoch = epoch
train_loss, train_perf = trainer.optimize(train_loader)
training_perf.append(train_perf)
# evaluate on validation set
if args.eval:
logging.info(f'Running validation on dev set')
val_loss, val_perf = trainer.evaluate(val_loader)
# remember best prec@1 and save checkpoint
if args.rank == 0:
is_best = val_loss < best_loss
best_loss = min(val_loss, best_loss)
trainer.save(save_all=args.save_all, is_best=is_best)
if args.eval:
utils.barrier()
eval_fname = f'eval_epoch_{epoch}'
eval_path = os.path.join(args.save_dir, eval_fname)
_, eval_stats = translator.run(
calc_bleu=True,
epoch=epoch,
eval_path=eval_path,
)
test_bleu = eval_stats['bleu']
if args.target_bleu and test_bleu >= args.target_bleu:
logging.info(f'Target accuracy reached')
break_training = True
acc_log = []
acc_log += [f'Summary: Epoch: {epoch}']
acc_log += [f'Training Loss: {train_loss:.4f}']
if args.eval:
acc_log += [f'Validation Loss: {val_loss:.4f}']
acc_log += [f'Test BLEU: {test_bleu:.2f}']
perf_log = []
perf_log += [f'Performance: Epoch: {epoch}']
perf_log += [f'Training: {train_perf:.0f} Tok/s']
if args.eval:
perf_log += [f'Validation: {val_perf:.0f} Tok/s']
if args.rank == 0:
logging.info('\t'.join(acc_log))
logging.info('\t'.join(perf_log))
logging.info(f'Finished epoch {epoch}')
if break_training:
break
utils.barrier()
training_stop = time.time()
training_time = training_stop - training_start
logging.info(f'Total training time {training_time:.0f} s')
table = TrainingTable()
avg_training_perf = sum(training_perf) / len(training_perf)
table.add(utils.get_world_size(), args.train_batch_size, test_bleu,
avg_training_perf, training_time)
if utils.get_rank() == 0:
table.write('Training Summary', args.math)
passed = utils.benchmark(test_bleu, args.target_bleu, train_perf,
args.target_perf)
if not passed:
sys.exit(1)
def main():
args = parse_args()
init_distributed(args)
if args.local_rank == 0:
dllogger.init(backends=[dllogger.JSONStreamBackend(verbosity=dllogger.Verbosity.VERBOSE,
filename=args.log_path),
dllogger.StdOutBackend(verbosity=dllogger.Verbosity.VERBOSE)])
else:
dllogger.init(backends=[])
dllogger.log(data=vars(args), step='PARAMETER')
if args.seed is not None:
torch.manual_seed(args.seed)
print("Saving results to {}".format(args.checkpoint_dir))
if not os.path.exists(args.checkpoint_dir) and args.checkpoint_dir != '':
os.makedirs(args.checkpoint_dir, exist_ok=True)
# sync workers before timing
if args.distributed:
torch.distributed.broadcast(torch.tensor([1], device="cuda"), 0)
torch.cuda.synchronize()
main_start_time = time.time()
train_ratings = torch.load(args.data+'/train_ratings.pt', map_location=torch.device('cuda:{}'.format(args.local_rank)))
test_ratings = torch.load(args.data+'/test_ratings.pt', map_location=torch.device('cuda:{}'.format(args.local_rank)))
test_negs = torch.load(args.data+'/test_negatives.pt', map_location=torch.device('cuda:{}'.format(args.local_rank)))
valid_negative = test_negs.shape[1]
nb_maxs = torch.max(train_ratings, 0)[0]
nb_users = nb_maxs[0].item() + 1
nb_items = nb_maxs[1].item() + 1
all_test_users = test_ratings.shape[0]
test_users, test_items, dup_mask, real_indices = dataloading.create_test_data(test_ratings, test_negs, args)
# make pytorch memory behavior more consistent later
torch.cuda.empty_cache()
# Create model
model = NeuMF(nb_users, nb_items,
mf_dim=args.factors,
mlp_layer_sizes=args.layers,
dropout=args.dropout)
optimizer = FusedAdam(model.parameters(), lr=args.learning_rate,
betas=(args.beta1, args.beta2), eps=args.eps)
criterion = nn.BCEWithLogitsLoss(reduction='none') # use torch.mean() with dim later to avoid copy to host
# Move model and loss to GPU
model = model.cuda()
criterion = criterion.cuda()
local_batch = args.batch_size // args.world_size
#traced_criterion = torch.jit.trace(criterion.forward,
# (torch.rand(local_batch,1),torch.rand(local_batch,1)))
traced_criterion = criterion
pyprof.init()
#import importlib
#pyprof.wrap(importlib.import_module(__name__), "traced_criterion")
#pyprof.wrap(traced_criterion, "__call__")
if args.opt_level == "O2":
model, optimizer = amp.initialize(model, optimizer, opt_level=args.opt_level,
keep_batchnorm_fp32=False, loss_scale='dynamic')
if args.distributed:
model = DDP(model)
print(model)
print("{} parameters".format(utils.count_parameters(model)))
if args.load_checkpoint_path:
state_dict = torch.load(args.load_checkpoint_path)
state_dict = {k.replace('module.', '') : v for k,v in state_dict.items()}
model.load_state_dict(state_dict)
if args.mode == 'test':
start = time.time()
hr, ndcg = val_epoch(model, test_users, test_items, dup_mask, real_indices, args.topk,
samples_per_user=valid_negative + 1,
num_user=all_test_users, distributed=args.distributed)
val_time = time.time() - start
eval_size = all_test_users * (valid_negative + 1)
eval_throughput = eval_size / val_time
dllogger.log(step=tuple(), data={'best_eval_throughput' : eval_throughput,
'hr@10' : hr})
return
max_hr = 0
best_epoch = 0
train_throughputs, eval_throughputs = [], []
with torch.autograd.profiler.emit_nvtx():
for epoch in range(args.epochs):
begin = time.time()
epoch_users, epoch_items, epoch_label = dataloading.prepare_epoch_train_data(train_ratings, nb_items, args)
num_batches = len(epoch_users)
for i in range(num_batches // args.grads_accumulated):
if i == 10:
profiler.start()
for j in range(args.grads_accumulated):
batch_idx = (args.grads_accumulated * i) + j
user = epoch_users[batch_idx]
item = epoch_items[batch_idx]
label = epoch_label[batch_idx].view(-1,1)
outputs = model(user, item)
nvtx.range_push("layer:Loss")
loss = traced_criterion(outputs, label).float()
nvtx.range_pop()
nvtx.range_push("layer:Mean")
loss = torch.mean(loss.view(-1), 0)
nvtx.range_pop()
if args.opt_level == "O2":
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
nvtx.range_push("layer:Adam")
optimizer.step()
nvtx.range_pop()
if i == 10:
profiler.stop()
for p in model.parameters():
p.grad = None
del epoch_users, epoch_items, epoch_label
train_time = time.time() - begin
begin = time.time()
epoch_samples = len(train_ratings) * (args.negative_samples + 1)
train_throughput = epoch_samples / train_time
train_throughputs.append(train_throughput)
hr, ndcg = val_epoch(model, test_users, test_items, dup_mask, real_indices, args.topk,
samples_per_user=valid_negative + 1,
num_user=all_test_users, epoch=epoch, distributed=args.distributed)
val_time = time.time() - begin
eval_size = all_test_users * (valid_negative + 1)
eval_throughput = eval_size / val_time
eval_throughputs.append(eval_throughput)
dllogger.log(step=(epoch,),
data = {'train_throughput': train_throughput,
'hr@10': hr,
'train_epoch_time': train_time,
'validation_epoch_time': val_time,
'eval_throughput': eval_throughput})
if hr > max_hr and args.local_rank == 0:
max_hr = hr
best_epoch = epoch
save_checkpoint_path = os.path.join(args.checkpoint_dir, 'model.pth')
print("New best hr! Saving the model to: ", save_checkpoint_path)
torch.save(model.state_dict(), save_checkpoint_path)
best_model_timestamp = time.time()
if args.threshold is not None:
if hr >= args.threshold:
print("Hit threshold of {}".format(args.threshold))
break
if args.local_rank == 0:
dllogger.log(data={'best_train_throughput': max(train_throughputs),
'best_eval_throughput': max(eval_throughputs),
'mean_train_throughput': np.mean(train_throughputs),
'mean_eval_throughput': np.mean(eval_throughputs),
'best_accuracy': max_hr,
'best_epoch': best_epoch,
'time_to_target': time.time() - main_start_time,
'time_to_best_model': best_model_timestamp - main_start_time},
step=tuple())