def train300_mlperf_coco(args):
args = setup_distributed(args)
# Build the model
model_options = {
'use_nhwc': args.nhwc,
'pad_input': args.pad_input,
'bn_group': args.bn_group,
}
ssd300 = SSD300(args, args.num_classes, **model_options)
if args.checkpoint is not None:
load_checkpoint(ssd300, args.checkpoint)
ssd300.train()
ssd300.cuda()
dboxes = dboxes300_coco()
# Note: No reason not to use optimised loss
loss_func = OptLoss()
loss_func.cuda()
# Create optimizer. This must also be done after network_to_half.
global_batch_size = (args.N_gpu * args.batch_size)
log_event(key=constants.MODEL_BN_SPAN,
value=args.bn_group * args.batch_size)
log_event(key=constants.GLOBAL_BATCH_SIZE, value=global_batch_size)
# mlperf only allows base_lr scaled by an integer
base_lr = 2.5e-3
requested_lr_multiplier = args.lr / base_lr
adjusted_multiplier = max(
1, round(requested_lr_multiplier * global_batch_size / 32))
current_lr = base_lr * adjusted_multiplier
current_momentum = 0.9
current_weight_decay = args.wd
static_loss_scale = 128.
optim = apex.optimizers.FusedSGD(ssd300.parameters(),
lr=current_lr,
momentum=current_momentum,
weight_decay=current_weight_decay)
ssd300, optim = apex.amp.initialize(ssd300,
optim,
opt_level='O2',
loss_scale=static_loss_scale)
# Parallelize. Need to do this after network_to_half.
if args.distributed:
if args.delay_allreduce:
print_message(args.local_rank,
"Delaying allreduces to the end of backward()")
ssd300 = DDP(ssd300,
gradient_predivide_factor=args.N_gpu / 8.0,
delay_allreduce=args.delay_allreduce,
retain_allreduce_buffers=args.use_fp16)
log_event(key=constants.OPT_BASE_LR, value=current_lr)
log_event(key=constants.OPT_LR_DECAY_BOUNDARY_EPOCHS,
value=args.lr_decay_epochs)
log_event(key=constants.OPT_LR_DECAY_STEPS, value=args.lr_decay_epochs)
log_event(key=constants.OPT_WEIGHT_DECAY, value=current_weight_decay)
if args.warmup is not None:
log_event(key=constants.OPT_LR_WARMUP_STEPS, value=args.warmup)
log_event(key=constants.OPT_LR_WARMUP_FACTOR, value=args.warmup_factor)
# Model is completely finished -- need to create separate copies, preserve parameters across
# them, and jit
ssd300_eval = SSD300(args, args.num_classes, **model_options).cuda()
if args.use_fp16:
convert_network(ssd300_eval, torch.half)
# Get the existant state from the train model
# * if we use distributed, then we want .module
train_model = ssd300.module if args.distributed else ssd300
ssd300_eval.load_state_dict(train_model.state_dict())
ssd300_eval.eval()
print_message(args.local_rank, "epoch", "nbatch", "loss")
iter_num = args.iteration
avg_loss = 0.0
start_elapsed_time = time.time()
last_printed_iter = args.iteration
num_elapsed_samples = 0
input_c = 4 if args.pad_input else 3
example_shape = [args.batch_size, 300, 300, input_c
] if args.nhwc else [args.batch_size, input_c, 300, 300]
example_input = torch.randn(*example_shape).cuda()
if args.use_fp16:
example_input = example_input.half()
if args.jit:
# DDP has some Python-side control flow. If we JIT the entire DDP-wrapped module,
# the resulting ScriptModule will elide this control flow, resulting in allreduce
# hooks not being called. If we're running distributed, we need to extract and JIT
# the wrapped .module.
# Replacing a DDP-ed ssd300 with a script_module might also cause the AccumulateGrad hooks
# to go out of scope, and therefore silently disappear.
module_to_jit = ssd300.module if args.distributed else ssd300
if args.distributed:
ssd300.module = torch.jit.trace(module_to_jit,
example_input,
check_trace=False)
else:
ssd300 = torch.jit.trace(module_to_jit,
example_input,
check_trace=False)
# JIT the eval model too
ssd300_eval = torch.jit.trace(ssd300_eval,
example_input,
check_trace=False)
# do a dummy fprop & bprop to make sure cudnnFind etc. are timed here
ploc, plabel = ssd300(example_input)
# produce a single dummy "loss" to make things easier
loss = ploc[0, 0, 0] + plabel[0, 0, 0]
dloss = torch.randn_like(loss)
# Cause cudnnFind for dgrad, wgrad to run
loss.backward(dloss)
# Necessary import in init
from pycocotools.coco import COCO
encoder = build_ssd300_coder()
evaluator = AsyncEvaluator(num_threads=1)
log_end(key=constants.INIT_STOP)
##### END INIT
# This is the first place we touch anything related to data
##### START DATA TOUCHING
barrier()
log_start(key=constants.RUN_START)
barrier()
train_pipe = prebuild_pipeline(args)
train_loader, epoch_size = build_pipeline(args,
training=True,
pipe=train_pipe)
if args.rank == 0:
print("epoch size is: ", epoch_size, " images")
val_loader, inv_map, cocoGt = build_pipeline(args, training=False)
if args.profile_gc_off:
gc.disable()
gc.collect()
##### END DATA TOUCHING
i_eval = 0
block_start_epoch = 1
log_start(key=constants.BLOCK_START,
metadata={
'first_epoch_num': block_start_epoch,
'epoch_count': args.evaluation[i_eval]
})
for epoch in range(args.epochs):
for p in ssd300.parameters():
p.grad = None
if epoch in args.evaluation:
# Get the existant state from the train model
# * if we use distributed, then we want .module
train_model = ssd300.module if args.distributed else ssd300
if args.distributed and args.allreduce_running_stats:
if args.rank == 0: print("averaging bn running means and vars")
# make sure every node has the same running bn stats before
# using them to evaluate, or saving the model for inference
world_size = float(torch.distributed.get_world_size())
for bn_name, bn_buf in train_model.named_buffers(recurse=True):
if ('running_mean' in bn_name) or ('running_var'
in bn_name):
torch.distributed.all_reduce(bn_buf,
op=dist.ReduceOp.SUM)
bn_buf /= world_size
if args.rank == 0:
if args.save:
print("saving model...")
if not os.path.isdir('./models'):
os.mkdir('./models')
torch.save({"model": ssd300.state_dict()},
"./models/iter_{}.pt".format(iter_num))
ssd300_eval.load_state_dict(train_model.state_dict())
# Note: No longer returns, evaluation is abstracted away inside evaluator
coco_eval(args,
ssd300_eval,
val_loader,
cocoGt,
encoder,
inv_map,
epoch,
iter_num,
evaluator=evaluator)
log_end(key=constants.BLOCK_STOP,
metadata={'first_epoch_num': block_start_epoch})
if epoch != max(args.evaluation):
i_eval += 1
block_start_epoch = epoch + 1
log_start(key=constants.BLOCK_START,
metadata={
'first_epoch_num':
block_start_epoch,
'epoch_count': (args.evaluation[i_eval] -
args.evaluation[i_eval - 1])
})
if epoch in args.lr_decay_epochs:
current_lr *= args.lr_decay_factor
print_message(
args.rank,
"lr decay step #" + str(bisect(args.lr_decay_epochs, epoch)))
for param_group in optim.param_groups:
param_group['lr'] = current_lr
log_start(key=constants.EPOCH_START,
metadata={
'epoch_num': epoch + 1,
'current_iter_num': iter_num
})
for i, (img, bbox, label) in enumerate(train_loader):
if args.profile_start is not None and iter_num == args.profile_start:
torch.cuda.profiler.start()
torch.cuda.synchronize()
if args.profile_nvtx:
torch.autograd._enable_profiler(
torch.autograd.ProfilerState.NVTX)
if args.profile is not None and iter_num == args.profile:
if args.profile_start is not None and iter_num >= args.profile_start:
# we turned cuda and nvtx profiling on, better turn it off too
if args.profile_nvtx:
torch.autograd._disable_profiler()
torch.cuda.profiler.stop()
return
if args.warmup is not None:
lr_warmup(optim, args.warmup, iter_num, epoch, current_lr,
args)
if (img is None) or (bbox is None) or (label is None):
print("No labels in batch")
continue
ploc, plabel = ssd300(img)
ploc, plabel = ploc.float(), plabel.float()
N = img.shape[0]
bbox.requires_grad = False
label.requires_grad = False
# reshape (N*8732X4 -> Nx8732x4) and transpose (Nx8732x4 -> Nx4x8732)
bbox = bbox.view(N, -1, 4).transpose(1, 2).contiguous()
# reshape (N*8732 -> Nx8732) and cast to Long
label = label.view(N, -1).long()
loss = loss_func(ploc, plabel, bbox, label)
if np.isfinite(loss.item()):
avg_loss = 0.999 * avg_loss + 0.001 * loss.item()
else:
print("model exploded (corrupted by Inf or Nan)")
sys.exit()
num_elapsed_samples += N
if args.rank == 0 and iter_num % args.print_interval == 0:
end_elapsed_time = time.time()
elapsed_time = end_elapsed_time - start_elapsed_time
avg_samples_per_sec = num_elapsed_samples * args.N_gpu / elapsed_time
print("Iteration: {:6d}, Loss function: {:5.3f}, Average Loss: {:.3f}, avg. samples / sec: {:.2f}"\
.format(iter_num, loss.item(), avg_loss, avg_samples_per_sec), end="\n")
last_printed_iter = iter_num
start_elapsed_time = time.time()
num_elapsed_samples = 0
with apex.amp.scale_loss(loss, optim) as scaled_loss:
scaled_loss.backward()
if not args.profile_fake_optim:
optim.step()
# Likely a decent skew here, let's take this opportunity to set the
# gradients to None. After DALI integration, playing with the
# placement of this is worth trying.
for p in ssd300.parameters():
p.grad = None
# Don't check every iteration due to cost of broadcast
if iter_num % 20 == 0:
finished = check_async_evals(args, evaluator, args.threshold)
if finished:
return True
iter_num += 1
train_loader.reset()
log_end(key=constants.EPOCH_STOP, metadata={'epoch_num': epoch + 1})
return False
def train300_mlperf_coco(args):
from pycocotools.coco import COCO
# Check that GPUs are actually available
use_cuda = not args.no_cuda
# Setup multi-GPU if necessary
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
if args.distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
local_seed = set_seeds(args)
# start timing here
if args.distributed:
N_gpu = torch.distributed.get_world_size()
else:
N_gpu = 1
validate_group_bn(args.bn_group)
# Setup data, defaults
dboxes = dboxes300_coco()
encoder = Encoder(dboxes)
input_size = 300
val_trans = SSDTransformer(dboxes, (input_size, input_size), val=True)
val_annotate = os.path.join(args.data,
"annotations/instances_val2017.json")
val_coco_root = os.path.join(args.data, "val2017")
train_annotate = os.path.join(args.data,
"annotations/instances_train2017.json")
train_coco_root = os.path.join(args.data, "train2017")
# Build the model
model_options = {
'backbone': args.backbone,
'use_nhwc': args.nhwc,
'pad_input': args.pad_input,
'bn_group': args.bn_group,
}
ssd300 = SSD300(args.num_classes, **model_options)
if args.checkpoint is not None:
load_checkpoint(ssd300, args.checkpoint)
ssd300.train()
ssd300.cuda()
if args.opt_loss:
loss_func = OptLoss(dboxes)
else:
loss_func = Loss(dboxes)
loss_func.cuda()
if args.distributed:
N_gpu = torch.distributed.get_world_size()
else:
N_gpu = 1
if args.use_fp16:
ssd300 = network_to_half(ssd300)
# Parallelize. Need to do this after network_to_half.
if args.distributed:
if args.delay_allreduce:
print_message(args.local_rank,
"Delaying allreduces to the end of backward()")
ssd300 = DDP(ssd300,
gradient_predivide_factor=N_gpu / 8.0,
delay_allreduce=args.delay_allreduce,
retain_allreduce_buffers=args.use_fp16)
# Create optimizer. This must also be done after network_to_half.
global_batch_size = (N_gpu * args.batch_size)
mlperf_print(key=mlperf_compliance.constants.MODEL_BN_SPAN,
value=args.bn_group * args.batch_size)
mlperf_print(key=mlperf_compliance.constants.GLOBAL_BATCH_SIZE,
value=global_batch_size)
# mlperf only allows base_lr scaled by an integer
base_lr = 2.5e-3
requested_lr_multiplier = args.lr / base_lr
adjusted_multiplier = max(
1, round(requested_lr_multiplier * global_batch_size / 32))
current_lr = base_lr * adjusted_multiplier
current_momentum = 0.9
current_weight_decay = args.wd
static_loss_scale = 128.
if args.use_fp16:
if args.distributed and not args.delay_allreduce:
# We can't create the flat master params yet, because we need to
# imitate the flattened bucket structure that DDP produces.
optimizer_created = False
else:
model_buckets = [
[
p for p in ssd300.parameters()
if p.requires_grad and p.type() == "torch.cuda.HalfTensor"
],
[
p for p in ssd300.parameters()
if p.requires_grad and p.type() == "torch.cuda.FloatTensor"
]
]
flat_master_buckets = create_flat_master(model_buckets)
optim = torch.optim.SGD(flat_master_buckets,
lr=current_lr,
momentum=current_momentum,
weight_decay=current_weight_decay)
optimizer_created = True
else:
optim = torch.optim.SGD(ssd300.parameters(),
lr=current_lr,
momentum=current_momentum,
weight_decay=current_weight_decay)
optimizer_created = True
mlperf_print(key=mlperf_compliance.constants.OPT_BASE_LR, value=current_lr)
mlperf_print(key=mlperf_compliance.constants.OPT_WEIGHT_DECAY,
value=current_weight_decay)
if args.warmup is not None:
mlperf_print(key=mlperf_compliance.constants.OPT_LR_WARMUP_STEPS,
value=args.warmup)
mlperf_print(key=mlperf_compliance.constants.OPT_LR_WARMUP_FACTOR,
value=args.warmup_factor)
# Model is completely finished -- need to create separate copies, preserve parameters across
# them, and jit
ssd300_eval = SSD300(args.num_classes,
backbone=args.backbone,
use_nhwc=args.nhwc,
pad_input=args.pad_input).cuda()
if args.use_fp16:
ssd300_eval = network_to_half(ssd300_eval)
# Get the existant state from the train model
# * if we use distributed, then we want .module
train_model = ssd300.module if args.distributed else ssd300
ssd300_eval.load_state_dict(train_model.state_dict())
ssd300_eval.eval()
print_message(args.local_rank, "epoch", "nbatch", "loss")
eval_points = np.array(args.evaluation) * 32 / global_batch_size
eval_points = list(map(int, list(eval_points)))
iter_num = args.iteration
avg_loss = 0.0
start_elapsed_time = time.time()
last_printed_iter = args.iteration
num_elapsed_samples = 0
# Generate normalization tensors
mean, std = generate_mean_std(args)
dummy_overflow_buf = torch.cuda.IntTensor([0])
def step_maybe_fp16_maybe_distributed(optim):
if args.use_fp16:
if args.distributed:
for flat_master, allreduce_buffer in zip(
flat_master_buckets, ssd300.allreduce_buffers):
if allreduce_buffer is None:
raise RuntimeError("allreduce_buffer is None")
flat_master.grad = allreduce_buffer.float()
flat_master.grad.data.mul_(1. / static_loss_scale)
else:
for flat_master, model_bucket in zip(flat_master_buckets,
model_buckets):
flat_grad = apex_C.flatten(
[m.grad.data for m in model_bucket])
flat_master.grad = flat_grad.float()
flat_master.grad.data.mul_(1. / static_loss_scale)
optim.step()
if args.use_fp16:
# Use multi-tensor scale instead of loop & individual parameter copies
for model_bucket, flat_master in zip(model_buckets,
flat_master_buckets):
multi_tensor_applier(
amp_C.multi_tensor_scale, dummy_overflow_buf, [
apex_C.unflatten(flat_master.data, model_bucket),
model_bucket
], 1.0)
input_c = 4 if args.pad_input else 3
example_shape = [args.batch_size, 300, 300, input_c
] if args.nhwc else [args.batch_size, input_c, 300, 300]
example_input = torch.randn(*example_shape).cuda()
if args.use_fp16:
example_input = example_input.half()
if args.jit:
# DDP has some Python-side control flow. If we JIT the entire DDP-wrapped module,
# the resulting ScriptModule will elide this control flow, resulting in allreduce
# hooks not being called. If we're running distributed, we need to extract and JIT
# the wrapped .module.
# Replacing a DDP-ed ssd300 with a script_module might also cause the AccumulateGrad hooks
# to go out of scope, and therefore silently disappear.
module_to_jit = ssd300.module if args.distributed else ssd300
if args.distributed:
ssd300.module = torch.jit.trace(module_to_jit, example_input)
else:
ssd300 = torch.jit.trace(module_to_jit, example_input)
# JIT the eval model too
ssd300_eval = torch.jit.trace(ssd300_eval, example_input)
# do a dummy fprop & bprop to make sure cudnnFind etc. are timed here
ploc, plabel = ssd300(example_input)
# produce a single dummy "loss" to make things easier
loss = ploc[0, 0, 0] + plabel[0, 0, 0]
dloss = torch.randn_like(loss)
# Cause cudnnFind for dgrad, wgrad to run
loss.backward(dloss)
mlperf_print(key=mlperf_compliance.constants.INIT_STOP, sync=True)
##### END INIT
# This is the first place we touch anything related to data
##### START DATA TOUCHING
mlperf_print(key=mlperf_compliance.constants.RUN_START, sync=True)
barrier()
cocoGt = COCO(annotation_file=val_annotate, use_ext=True)
val_coco = COCODetection(val_coco_root, val_annotate, val_trans)
if args.distributed:
val_sampler = GeneralDistributedSampler(val_coco, pad=False)
else:
val_sampler = None
if args.no_dali:
train_trans = SSDTransformer(dboxes, (input_size, input_size),
val=False)
train_coco = COCODetection(train_coco_root, train_annotate,
train_trans)
if args.distributed:
train_sampler = GeneralDistributedSampler(train_coco, pad=False)
else:
train_sampler = None
train_loader = DataLoader(train_coco,
batch_size=args.batch_size *
args.input_batch_multiplier,
shuffle=(train_sampler is None),
sampler=train_sampler,
num_workers=args.num_workers,
collate_fn=partial(my_collate,
is_training=True))
else:
train_pipe = COCOPipeline(args.batch_size *
args.input_batch_multiplier,
args.local_rank,
train_coco_root,
train_annotate,
N_gpu,
num_threads=args.num_workers,
output_fp16=args.use_fp16,
output_nhwc=args.nhwc,
pad_output=args.pad_input,
seed=local_seed - 2**31,
use_nvjpeg=args.use_nvjpeg,
use_roi=args.use_roi_decode,
dali_cache=args.dali_cache,
dali_async=(not args.dali_sync))
print_message(args.local_rank,
"time_check a: {secs:.9f}".format(secs=time.time()))
train_pipe.build()
print_message(args.local_rank,
"time_check b: {secs:.9f}".format(secs=time.time()))
test_run = train_pipe.run()
train_loader = SingleDaliIterator(
train_pipe, [
'images',
DALIOutput('bboxes', False, True),
DALIOutput('labels', True, True)
],
train_pipe.epoch_size()['train_reader'],
ngpu=N_gpu)
train_loader = EncodingInputIterator(train_loader,
dboxes=encoder.dboxes.cuda(),
nhwc=args.nhwc,
fake_input=args.fake_input,
no_dali=args.no_dali)
if args.input_batch_multiplier > 1:
train_loader = RateMatcher(input_it=train_loader,
output_size=args.batch_size)
val_dataloader = DataLoader(
val_coco,
batch_size=args.eval_batch_size,
shuffle=False, # Note: distributed sampler is shuffled :(
sampler=val_sampler,
num_workers=args.num_workers)
inv_map = {v: k for k, v in val_coco.label_map.items()}
##### END DATA TOUCHING
i_eval = 0
first_epoch = 1
mlperf_print(key=mlperf_compliance.constants.BLOCK_START,
metadata={
'first_epoch_num':
first_epoch,
'epoch_count':
args.evaluation[i_eval] * 32 /
train_pipe.epoch_size()['train_reader']
},
sync=True)
for epoch in range(args.epochs):
mlperf_print(key=mlperf_compliance.constants.EPOCH_START,
metadata={'epoch_num': epoch + 1},
sync=True)
for p in ssd300.parameters():
p.grad = None
for i, (img, bbox, label) in enumerate(train_loader):
if args.profile_start is not None and iter_num == args.profile_start:
torch.cuda.profiler.start()
torch.cuda.synchronize()
if args.profile_nvtx:
torch.autograd._enable_profiler(
torch.autograd.ProfilerState.NVTX)
if args.profile is not None and iter_num == args.profile:
if args.profile_start is not None and iter_num >= args.profile_start:
# we turned cuda and nvtx profiling on, better turn it off too
if args.profile_nvtx:
torch.autograd._disable_profiler()
torch.cuda.profiler.stop()
return
if args.warmup is not None and optimizer_created:
lr_warmup(optim, args.warmup, iter_num, epoch, current_lr,
args)
if iter_num == ((args.decay1 * 1000 * 32) // global_batch_size):
print_message(args.local_rank, "lr decay step #1")
current_lr *= 0.1
for param_group in optim.param_groups:
param_group['lr'] = current_lr
if iter_num == ((args.decay2 * 1000 * 32) // global_batch_size):
print_message(args.local_rank, "lr decay step #2")
current_lr *= 0.1
for param_group in optim.param_groups:
param_group['lr'] = current_lr
if (img is None) or (bbox is None) or (label is None):
print("No labels in batch")
continue
ploc, plabel = ssd300(img)
ploc, plabel = ploc.float(), plabel.float()
N = img.shape[0]
gloc, glabel = Variable(bbox, requires_grad=False), \
Variable(label, requires_grad=False)
loss = loss_func(ploc, plabel, gloc, glabel)
if np.isfinite(loss.item()):
avg_loss = 0.999 * avg_loss + 0.001 * loss.item()
else:
print("model exploded (corrupted by Inf or Nan)")
sys.exit()
num_elapsed_samples += N
if args.local_rank == 0 and iter_num % args.print_interval == 0:
end_elapsed_time = time.time()
elapsed_time = end_elapsed_time - start_elapsed_time
avg_samples_per_sec = num_elapsed_samples * N_gpu / elapsed_time
print("Iteration: {:6d}, Loss function: {:5.3f}, Average Loss: {:.3f}, avg. samples / sec: {:.2f}"\
.format(iter_num, loss.item(), avg_loss, avg_samples_per_sec), end="\n")
last_printed_iter = iter_num
start_elapsed_time = time.time()
num_elapsed_samples = 0
# loss scaling
if args.use_fp16:
loss = loss * static_loss_scale
loss.backward()
if not optimizer_created:
# Imitate the model bucket structure created by DDP.
# These will already be split by type (float or half).
model_buckets = []
for bucket in ssd300.active_i_buckets:
model_buckets.append([])
for active_i in bucket:
model_buckets[-1].append(
ssd300.active_params[active_i])
flat_master_buckets = create_flat_master(model_buckets)
optim = torch.optim.SGD(flat_master_buckets,
lr=current_lr,
momentum=current_momentum,
weight_decay=current_weight_decay)
optimizer_created = True
# Skip this first iteration because flattened allreduce buffers are not yet created.
# step_maybe_fp16_maybe_distributed(optim)
else:
step_maybe_fp16_maybe_distributed(optim)
# Likely a decent skew here, let's take this opportunity to set the gradients to None.
# After DALI integration, playing with the placement of this is worth trying.
for p in ssd300.parameters():
p.grad = None
if iter_num in eval_points:
# Get the existant state from the train model
# * if we use distributed, then we want .module
train_model = ssd300.module if args.distributed else ssd300
if args.distributed and args.allreduce_running_stats:
if get_rank() == 0:
print("averaging bn running means and vars")
# make sure every node has the same running bn stats before
# using them to evaluate, or saving the model for inference
world_size = float(torch.distributed.get_world_size())
for bn_name, bn_buf in train_model.named_buffers(
recurse=True):
if ('running_mean' in bn_name) or ('running_var'
in bn_name):
torch.distributed.all_reduce(bn_buf,
op=dist.ReduceOp.SUM)
bn_buf /= world_size
if get_rank() == 0:
if not args.no_save:
print("saving model...")
torch.save(
{
"model": ssd300.state_dict(),
"label_map": val_coco.label_info
}, "./models/iter_{}.pt".format(iter_num))
ssd300_eval.load_state_dict(train_model.state_dict())
succ = coco_eval(
ssd300_eval,
val_dataloader,
cocoGt,
encoder,
inv_map,
args.threshold,
epoch,
iter_num,
args.eval_batch_size,
use_fp16=args.use_fp16,
local_rank=args.local_rank if args.distributed else -1,
N_gpu=N_gpu,
use_nhwc=args.nhwc,
pad_input=args.pad_input)
mlperf_print(key=mlperf_compliance.constants.BLOCK_STOP,
metadata={'first_epoch_num': first_epoch},
sync=True)
if succ:
return True
if iter_num != max(eval_points):
i_eval += 1
first_epoch = epoch + 1
mlperf_print(key=mlperf_compliance.constants.BLOCK_START,
metadata={
'first_epoch_num':
first_epoch,
'epoch_count':
(args.evaluation[i_eval] -
args.evaluation[i_eval - 1]) * 32 /
train_pipe.epoch_size()['train_reader']
},
sync=True)
iter_num += 1
if args.max_iter > 0:
if iter_num > args.max_iter:
break
train_loader.reset()
mlperf_print(key=mlperf_compliance.constants.EPOCH_STOP,
metadata={'epoch_num': epoch + 1},
sync=True)
return False
def main():
args = parse_arguments()
status = 'aborted' # later set to 'success' if termination criteria met
mlperf_logger.log_start(key=mlperf_logger.constants.INIT_START,
log_all_ranks=True, sync=False)
if args.use_env and 'LOCAL_RANK' in os.environ:
args.local_rank = int(os.environ['LOCAL_RANK'])
device, args = setup_training(args)
mlperf_logger.mlperf_submission_log('bert')
worker_seeds, shuffling_seeds = utils.setup_seeds(args.seed, args.num_epochs_to_generate_seeds_for, device)
worker_seed = worker_seeds[torch.distributed.get_rank()]
random.seed(worker_seed)
np.random.seed(worker_seed)
torch.manual_seed(worker_seed)
worker_init = WorkerInitObj(worker_seed)
mlperf_logger.log_event(key=mlperf_logger.constants.SEED, value=args.seed,
sync=False)
mlperf_logger.log_event(key=mlperf_logger.constants.GLOBAL_BATCH_SIZE,
value=global_batch_size(args), sync=False)
mlperf_logger.log_event(key='opt_gradient_accumulation_steps',
value=args.gradient_accumulation_steps, sync=False)
mlperf_logger.log_event(key='max_predictions_per_seq',
value=args.max_predictions_per_seq, sync=False)
mlperf_logger.log_event(key='opt_learning_rate_training_steps',
value=args.max_steps, sync=False)
mlperf_logger.log_event(key='num_warmup_steps',
value=int(args.warmup_proportion*args.max_steps) if args.warmup_steps==0 else args.warmup_steps,
sync=False)
if utils.is_main_process():
print("parsed args:")
print(args)
# Prepare optimizer
model, optimizer, lr_scheduler, checkpoint, global_step = prepare_model_and_optimizer(args, device)
samples_trained = global_step * args.train_batch_size * args.gradient_accumulation_steps * args.n_gpu
if args.unpad:
torch.cuda.synchronize()
InitMHACUDAExtension()
torch.cuda.synchronize()
final_loss = float("inf")
train_time_raw = float("inf")
raw_train_start = time.time()
if args.do_train:
model.train()
most_recent_ckpts_paths = []
average_loss = 0.0 # averaged loss every args.log_freq steps
epoch = 1
training_steps = 0
end_training, converged = False, False
samples_trained_prev = 0
eval_count = 0
pool = ProcessPoolExecutor(1)
if args.target_mlm_accuracy:
if args.train_mlm_accuracy_window_size > 0:
accuracy_scores = []
avg_mlm_accuracy = torch.Tensor([0]).cuda()
first_epoch = True
if found_resume_checkpoint(args):
f_start_id = checkpoint['files'][0]
files = checkpoint['files'][1:]
num_files = len(files)
else:
files = [os.path.join(args.input_dir, f) for f in os.listdir(args.input_dir) if
os.path.isfile(os.path.join(args.input_dir, f)) and 'part' in f]
files.sort()
num_files = len(files)
random.Random(shuffling_seeds[epoch]).shuffle(files)
f_start_id = 0
mlperf_logger.log_end(key=mlperf_logger.constants.INIT_STOP, sync=False)
mlperf_logger.log_start(key=mlperf_logger.constants.RUN_START, sync=True)
mlperf_logger.barrier()
# Start prefetching eval dataset
if args.eval_dir:
eval_dataset_future = pool.submit(create_eval_dataset, args, worker_init_fn=worker_init)
while global_step < args.max_steps and not end_training:
mlperf_logger.log_start(key=mlperf_logger.constants.EPOCH_START,
metadata={'epoch_num': epoch}, sync=False)
mlperf_logger.log_start(key=mlperf_logger.constants.BLOCK_START,
metadata={'first_epoch_num': epoch,
'epoch_count': 1},
sync=False)
if utils.is_main_process():
print("parsed args:")
print(args)
now_time = time.time()
now_step = global_step
now_skipped = skipped_steps
print("epoch:", epoch)
thread = None
# Reshuffle file list on subsequent epochs
if not first_epoch:
files = [os.path.join(args.input_dir, f) for f in os.listdir(args.input_dir) if
os.path.isfile(os.path.join(args.input_dir, f)) and 'part' in f]
files.sort()
num_files = len(files)
random.Random(shuffling_seeds[epoch]).shuffle(files)
f_start_id = 0
first_epoch = False
shared_file_list = {}
if torch.distributed.is_initialized() and torch.distributed.get_world_size() > num_files:
remainder = torch.distributed.get_world_size() % num_files
data_file = files[(f_start_id*torch.distributed.get_world_size() + torch.distributed.get_rank() +
remainder * f_start_id) % num_files]
else:
data_file = files[(f_start_id*torch.distributed.get_world_size() + torch.distributed.get_rank()) % num_files]
previous_file = data_file
train_data = pretraining_dataset(data_file, args.max_predictions_per_seq)
train_sampler = RandomSampler(train_data)
train_dataloader = DataLoader(train_data, sampler=train_sampler,
batch_size=args.train_batch_size, num_workers=4, worker_init_fn=worker_init, pin_memory=True)
overflow_buf = None
if args.allreduce_post_accumulation:
overflow_buf = torch.cuda.IntTensor([0])
for f_id in range(f_start_id + 1, len(files)):
if torch.distributed.get_world_size() > num_files:
data_file = files[(f_id*torch.distributed.get_world_size() + torch.distributed.get_rank() +
remainder * f_id) % num_files]
else:
data_file = files[(f_id*torch.distributed.get_world_size() + torch.distributed.get_rank())%num_files]
previous_file = data_file
dataset_future = pool.submit(create_pretraining_dataset, data_file, args.max_predictions_per_seq, shared_file_list, args, worker_init_fn=worker_init)
for step, batch in enumerate(train_dataloader):
training_steps += 1
update_step = training_steps % args.gradient_accumulation_steps == 0
batch = [t.to(device) for t in batch]
input_ids, segment_ids, input_mask, masked_lm_labels, next_sentence_labels = batch
loss, mlm_acc, _ = model(input_ids=input_ids, token_type_ids=segment_ids, attention_mask=input_mask,
masked_lm_labels=masked_lm_labels, next_sentence_label=next_sentence_labels,
checkpoint_activations=args.checkpoint_activations)
divisor = args.gradient_accumulation_steps
if args.gradient_accumulation_steps > 1:
if not args.allreduce_post_accumulation:
# this division was merged into predivision
loss = loss / args.gradient_accumulation_steps
divisor = 1.0
if args.fp16:
with amp.scale_loss(loss, optimizer, delay_overflow_check=args.allreduce_post_accumulation) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
average_loss += loss.item()
if update_step:
lr_scheduler.step() # learning rate warmup
global_step = take_optimizer_step(args, optimizer, model, overflow_buf, global_step)
samples_trained = global_step * args.train_batch_size * args.gradient_accumulation_steps * args.n_gpu
if (args.eval_dir and args.eval_iter_samples > 0 and
samples_trained >= args.eval_iter_start_samples + eval_count * args.eval_iter_samples):
# on first eval, get eval_dataloader
if eval_count == 0:
eval_dataloader = eval_dataset_future.result(timeout=None)
samples_trained_prev = samples_trained
eval_avg_loss, eval_avg_mlm_accuracy = run_eval(model, eval_dataloader, device, args.num_eval_examples,
first_eval=(eval_count == 0), use_cache=args.cache_eval_data)
if utils.is_main_process():
mlperf_logger.log_event(key=mlperf_logger.constants.EVAL_ACCURACY, value=eval_avg_mlm_accuracy, metadata={'epoch_num': epoch}, sync=False)
print({"global_steps": global_step, "eval_loss": eval_avg_loss, "eval_mlm_accuracy":eval_avg_mlm_accuracy})
if args.target_mlm_accuracy:
if eval_avg_mlm_accuracy >= args.target_mlm_accuracy:
end_training, converged = True, True
if utils.is_main_process():
print("%f > %f, Target MLM Accuracy reached at %d"%(eval_avg_mlm_accuracy, args.target_mlm_accuracy, global_step))
eval_count += 1
if args.target_mlm_accuracy and args.train_mlm_accuracy_window_size > 0:
accuracy_scores.append(mlm_acc)
if update_step:
accuracy_scores = accuracy_scores[-args.train_mlm_accuracy_window_size * args.gradient_accumulation_steps:]
avg_mlm_accuracy[0] = sum(accuracy_scores) / len(accuracy_scores)
torch.distributed.all_reduce(avg_mlm_accuracy, op=torch.distributed.ReduceOp.SUM)
avg_mlm_accuracy /= torch.distributed.get_world_size()
if training_steps % (args.log_freq * args.gradient_accumulation_steps) == 0:
samples_trained = global_step * args.train_batch_size * args.gradient_accumulation_steps * args.n_gpu
if utils.is_main_process():
time_interval = time.time() - now_time
step_interval = global_step - now_step
skip_interval = skipped_steps - now_skipped
now_time = time.time()
now_step = global_step
now_skipped = skipped_steps
training_perf = args.train_batch_size * args.gradient_accumulation_steps * args.n_gpu \
* (step_interval + skip_interval) / time_interval
if args.train_mlm_accuracy_window_size > 0:
print({"training_steps": training_steps,
"average_loss": average_loss / (args.log_freq * divisor),
"step_loss": loss.item() * args.gradient_accumulation_steps / divisor,
"learning_rate": optimizer.param_groups[0]['lr'],
"seq/s": training_perf,
"global_steps": now_step,
"samples_trained": samples_trained,
"skipped_steps": now_skipped,
"timestamp": now_time,
"mlm_accuracy": avg_mlm_accuracy[0].item()})
else:
print({"training_steps": training_steps,
"average_loss": average_loss / (args.log_freq * divisor),
"step_loss": loss.item() * args.gradient_accumulation_steps / divisor,
"learning_rate": optimizer.param_groups[0]['lr'],
"seq/s": training_perf,
"global_steps": now_step,
"samples_trained": samples_trained,
"skipped_steps": now_skipped,
"timestamp": now_time})
average_loss = 0
if global_step >= args.max_steps or end_training:
status = 'success' if converged else 'aborted'
end_training = True
train_time_raw = time.time() - raw_train_start
last_num_steps = int(training_steps / args.gradient_accumulation_steps) % args.log_freq
last_num_steps = args.log_freq if last_num_steps == 0 else last_num_steps
average_loss = torch.tensor(average_loss, dtype=torch.float32).cuda()
average_loss = average_loss / (last_num_steps * divisor)
if (torch.distributed.is_initialized()):
average_loss /= torch.distributed.get_world_size()
torch.distributed.all_reduce(average_loss)
final_loss = average_loss.item()
if utils.is_main_process():
if args.train_mlm_accuracy_window_size > 0:
print((epoch, training_steps / args.gradient_accumulation_steps, ), {"final_loss": final_loss,
"final_mlm_accuracy": avg_mlm_accuracy[0].item()})
else:
print((epoch, training_steps / args.gradient_accumulation_steps, ), {"final_loss": final_loss})
if end_training or (samples_trained - samples_trained_prev >= args.num_samples_per_checkpoint and samples_trained >= args.min_samples_to_start_checkpoints):
samples_trained_prev = samples_trained
if utils.is_main_process() and not args.skip_checkpoint:
# Save a trained model
model_to_save = model.module if hasattr(model,
'module') else model # Only save the model it-self
if args.phase2:
output_save_file = os.path.join(args.output_dir, "phase2_ckpt_{}.pt".format(samples_trained))
else:
output_save_file = os.path.join(args.output_dir, "phase1_ckpt_{}.pt".format(samples_trained))
if args.do_train:
torch.save({'model': model_to_save.state_dict(),
'optimizer': optimizer.state_dict(),
'master params': list(amp.master_params(optimizer)),
'files': [f_id] + files}, output_save_file)
most_recent_ckpts_paths.append(output_save_file)
if len(most_recent_ckpts_paths) > args.keep_n_most_recent_checkpoints:
ckpt_to_be_removed = most_recent_ckpts_paths.pop(0)
os.remove(ckpt_to_be_removed)
if samples_trained >= args.max_samples_termination or end_training:
status = 'success' if converged else 'aborted'
end_training = True
break
del train_dataloader
if samples_trained >= args.max_samples_termination or end_training:
status = 'success' if converged else 'aborted'
end_training = True
break
train_dataloader, data_file = dataset_future.result(timeout=None)
mlperf_logger.log_end(key=mlperf_logger.constants.BLOCK_STOP,
metadata={'first_epoch_num': epoch},
sync=False)
mlperf_logger.log_end(key=mlperf_logger.constants.EPOCH_STOP,
metadata={'epoch_num': epoch}, sync=False)
epoch += 1
mlperf_logger.log_event(key=mlperf_logger.constants.TRAIN_SAMPLES,
value=samples_trained,
sync=False)
mlperf_logger.log_event(key=mlperf_logger.constants.EVAL_SAMPLES,
value=args.num_eval_examples,
sync=False)
mlperf_logger.log_end(key=mlperf_logger.constants.RUN_STOP,
metadata={'status': status}, sync=False)
return args, final_loss, train_time_raw