def train(epoch, config, model, optimizer, scheduler, loss_func, train_loader,
logger, tensorboard_writer, tensorboard_writer2):
global global_step
logger.info(f'Train {epoch} {global_step}')
device = torch.device(config.device)
model.train()
loss_meter = AverageMeter()
acc1_meter = AverageMeter()
acc5_meter = AverageMeter()
start = time.time()
for step, (data, targets) in enumerate(train_loader):
step += 1
global_step += 1
if get_rank() == 0 and step == 1:
if config.tensorboard.train_images:
image = torchvision.utils.make_grid(data,
normalize=True,
scale_each=True)
tensorboard_writer.add_image('Train/Image', image, epoch)
data = data.to(device,
non_blocking=config.train.dataloader.non_blocking)
targets = send_targets_to_device(config, targets, device)
data_chunks, target_chunks = subdivide_batch(config, data, targets)
optimizer.zero_grad()
outputs = []
losses = []
for data_chunk, target_chunk in zip(data_chunks, target_chunks):
if config.augmentation.use_dual_cutout:
w = data_chunk.size(3) // 2
data1 = data_chunk[:, :, :, :w]
data2 = data_chunk[:, :, :, w:]
outputs1 = model(data1)
outputs2 = model(data2)
output_chunk = torch.cat(
(outputs1.unsqueeze(1), outputs2.unsqueeze(1)), dim=1)
else:
output_chunk = model(data_chunk)
outputs.append(output_chunk)
loss = loss_func(output_chunk, target_chunk)
if "CIFAR10_CM" in config.dataset.name: # Added by W210 Team
loss = loss.mean()
losses.append(loss)
with apex.amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
outputs = torch.cat(outputs)
if config.train.gradient_clip > 0:
torch.nn.utils.clip_grad_norm_(apex.amp.master_params(optimizer),
config.train.gradient_clip)
if config.train.subdivision > 1:
for param in model.parameters():
param.grad.data.div_(config.train.subdivision)
optimizer.step()
acc1, acc5 = compute_accuracy(config,
outputs,
targets,
augmentation=True,
topk=(1, 5))
loss = sum(losses)
if config.train.distributed:
loss_all_reduce = dist.all_reduce(loss,
op=dist.ReduceOp.SUM,
async_op=True)
acc1_all_reduce = dist.all_reduce(acc1,
op=dist.ReduceOp.SUM,
async_op=True)
acc5_all_reduce = dist.all_reduce(acc5,
op=dist.ReduceOp.SUM,
async_op=True)
loss_all_reduce.wait()
acc1_all_reduce.wait()
acc5_all_reduce.wait()
loss.div_(dist.get_world_size())
acc1.div_(dist.get_world_size())
acc5.div_(dist.get_world_size())
loss = loss.item()
acc1 = acc1.item()
acc5 = acc5.item()
num = data.size(0)
loss_meter.update(loss, num)
acc1_meter.update(acc1, num)
acc5_meter.update(acc5, num)
torch.cuda.synchronize()
if get_rank() == 0:
if step % config.train.log_period == 0 or step == len(
train_loader):
logger.info(
f'Epoch {epoch} '
f'Step {step}/{len(train_loader)} '
f'lr {scheduler.get_last_lr()[0]:.6f} '
f'loss {loss_meter.val:.4f} ({loss_meter.avg:.4f}) '
f'acc@1 {acc1_meter.val:.4f} ({acc1_meter.avg:.4f}) '
f'acc@5 {acc5_meter.val:.4f} ({acc5_meter.avg:.4f})')
tensorboard_writer2.add_scalar('Train/RunningLoss',
loss_meter.avg, global_step)
tensorboard_writer2.add_scalar('Train/RunningAcc1',
acc1_meter.avg, global_step)
tensorboard_writer2.add_scalar('Train/RunningAcc5',
acc5_meter.avg, global_step)
tensorboard_writer2.add_scalar('Train/RunningLearningRate',
scheduler.get_last_lr()[0],
global_step)
scheduler.step()
if get_rank() == 0:
elapsed = time.time() - start
logger.info(f'Elapsed {elapsed:.2f}')
tensorboard_writer.add_scalar('Train/Loss', loss_meter.avg, epoch)
tensorboard_writer.add_scalar('Train/Acc1', acc1_meter.avg, epoch)
tensorboard_writer.add_scalar('Train/Acc5', acc5_meter.avg, epoch)
tensorboard_writer.add_scalar('Train/Time', elapsed, epoch)
tensorboard_writer.add_scalar('Train/LearningRate',
scheduler.get_last_lr()[0], epoch)
def validate(epoch, config, model, loss_func, val_loader, logger,
tensorboard_writer):
logger.info(f'Val {epoch}')
device = torch.device(config.device)
model.eval()
loss_meter = AverageMeter()
acc1_meter = AverageMeter()
acc5_meter = AverageMeter()
start = time.time()
with torch.no_grad():
for step, (data, targets) in enumerate(val_loader):
if get_rank() == 0:
if config.tensorboard.val_images:
if epoch == 0 and step == 0:
image = torchvision.utils.make_grid(data,
normalize=True,
scale_each=True)
tensorboard_writer.add_image('Val/Image', image, epoch)
data = data.to(
device, non_blocking=config.validation.dataloader.non_blocking)
targets = targets.to(device)
outputs = model(data)
loss = loss_func(outputs, targets)
acc1, acc5 = compute_accuracy(config,
outputs,
targets,
augmentation=False,
topk=(1, 5))
if config.train.distributed:
loss_all_reduce = dist.all_reduce(loss,
op=dist.ReduceOp.SUM,
async_op=True)
acc1_all_reduce = dist.all_reduce(acc1,
op=dist.ReduceOp.SUM,
async_op=True)
acc5_all_reduce = dist.all_reduce(acc5,
op=dist.ReduceOp.SUM,
async_op=True)
loss_all_reduce.wait()
acc1_all_reduce.wait()
acc5_all_reduce.wait()
loss.div_(dist.get_world_size())
acc1.div_(dist.get_world_size())
acc5.div_(dist.get_world_size())
loss = loss.item()
acc1 = acc1.item()
acc5 = acc5.item()
num = data.size(0)
loss_meter.update(loss, num)
acc1_meter.update(acc1, num)
acc5_meter.update(acc5, num)
torch.cuda.synchronize()
logger.info(f'Epoch {epoch} '
f'loss {loss_meter.avg:.4f} '
f'acc@1 {acc1_meter.avg:.4f} '
f'acc@5 {acc5_meter.avg:.4f}')
elapsed = time.time() - start
logger.info(f'Elapsed {elapsed:.2f}')
if get_rank() == 0:
if epoch > 0:
tensorboard_writer.add_scalar('Val/Loss', loss_meter.avg, epoch)
tensorboard_writer.add_scalar('Val/Acc1', acc1_meter.avg, epoch)
tensorboard_writer.add_scalar('Val/Acc5', acc5_meter.avg, epoch)
tensorboard_writer.add_scalar('Val/Time', elapsed, epoch)
if config.tensorboard.model_params:
for name, param in model.named_parameters():
tensorboard_writer.add_histogram(name, param, epoch)
def train(epoch, config, model, optimizer, scheduler, loss_func, train_loader,
logger, tensorboard_writer, tensorboard_writer2):
global global_step
logger.info(f'Train {epoch} {global_step}')
device = torch.device(config.device)
model.train()
loss_meter = AverageMeter()
acc1_meter = AverageMeter()
acc5_meter = AverageMeter()
start = time.time()
for step, (data, targets) in enumerate(train_loader):
# every step is an iteration
step += 1
global_step += 1
if get_rank() == 0 and step == 1:
if config.tensorboard.train_images:
image = torchvision.utils.make_grid(data,
normalize=True,
scale_each=True)
tensorboard_writer.add_image('Train/Image', image, epoch)
data = data.to(device,
non_blocking=config.train.dataloader.non_blocking)
# Because target is not also pure single target(label), when data augmentation like mixup is deployed,
# multiple labels could occur and need to be sent to device separately.
targets = send_targets_to_device(config, targets, device)
data_chunks, target_chunks = subdivide_batch(config, data, targets)
optimizer.zero_grad()
outputs = []
losses = []
for data_chunk, target_chunk in zip(data_chunks, target_chunks):
if config.augmentation.use_dual_cutout:
w = data_chunk.size(3) // 2
data1 = data_chunk[:, :, :, :w]
data2 = data_chunk[:, :, :, w:]
outputs1 = model(data1)
outputs2 = model(data2)
output_chunk = torch.cat(
(outputs1.unsqueeze(1), outputs2.unsqueeze(1)), dim=1)
else:
output_chunk = model(data_chunk)
outputs.append(output_chunk)
loss = loss_func(output_chunk, target_chunk)
# Loss is used for calculating and accumulating the gradients.
# But losses is a list containing all losses but not for gradients.
losses.append(loss)
if config.device != 'cpu':
with apex.amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
# concatenating all chunks into one piece.
outputs = torch.cat(outputs)
if config.train.gradient_clip > 0:
# If norm of gradients higher than the specified value in the config,
# scale the gradient to the specified gradient norm value.
if config.device != 'cpu':
torch.nn.utils.clip_grad_norm_(
apex.amp.master_params(optimizer),
config.train.gradient_clip)
else:
torch.nn.utils.clip_grad_norm_(model.parameters(),
config.train.gradient_clip)
if config.train.subdivision > 1:
for param in model.parameters():
# the final gradients should be divided(averaged) by the number of the subdivision
param.grad.data.div_(config.train.subdivision)
# optimizing gradients
optimizer.step()
acc1, acc5 = compute_accuracy(config,
outputs,
targets,
augmentation=True,
topk=(1, 5))
loss = sum(losses)
if config.train.distributed:
loss_all_reduce = dist.all_reduce(loss,
op=dist.ReduceOp.SUM,
async_op=True)
acc1_all_reduce = dist.all_reduce(acc1,
op=dist.ReduceOp.SUM,
async_op=True)
acc5_all_reduce = dist.all_reduce(acc5,
op=dist.ReduceOp.SUM,
async_op=True)
loss_all_reduce.wait()
acc1_all_reduce.wait()
acc5_all_reduce.wait()
loss.div_(dist.get_world_size())
acc1.div_(dist.get_world_size())
acc5.div_(dist.get_world_size())
loss = loss.item()
acc1 = acc1.item()
acc5 = acc5.item()
num = data.size(0)
loss_meter.update(loss, num)
acc1_meter.update(acc1, num)
acc5_meter.update(acc5, num)
torch.cuda.synchronize()
if get_rank() == 0:
if step % config.train.log_period == 0 or step == len(
train_loader):
logger.info(
f'Epoch {epoch} '
f'Step {step}/{len(train_loader)} '
f'lr {scheduler.get_last_lr()[0]:.6f} '
f'loss {loss_meter.val:.4f} ({loss_meter.avg:.4f}) '
f'acc@1 {acc1_meter.val:.4f} ({acc1_meter.avg:.4f}) '
f'acc@5 {acc5_meter.val:.4f} ({acc5_meter.avg:.4f})')
tensorboard_writer2.add_scalar('Train/RunningLoss',
loss_meter.avg, global_step)
tensorboard_writer2.add_scalar('Train/RunningAcc1',
acc1_meter.avg, global_step)
tensorboard_writer2.add_scalar('Train/RunningAcc5',
acc5_meter.avg, global_step)
tensorboard_writer2.add_scalar('Train/RunningLearningRate',
scheduler.get_last_lr()[0],
global_step)
scheduler.step()
if get_rank() == 0:
elapsed = time.time() - start
logger.info(f'Elapsed {elapsed:.2f}')
tensorboard_writer.add_scalar('Train/Loss', loss_meter.avg, epoch)
tensorboard_writer.add_scalar('Train/Acc1', acc1_meter.avg, epoch)
tensorboard_writer.add_scalar('Train/Acc5', acc5_meter.avg, epoch)
tensorboard_writer.add_scalar('Train/Time', elapsed, epoch)
tensorboard_writer.add_scalar('Train/LearningRate',
scheduler.get_last_lr()[0], epoch)
