def forward(ctx, input, num_sync_devices, num_groups):
"""
Perform forwarding, gathering the stats across different process/ GPU
group.
"""
ctx.num_sync_devices = num_sync_devices
ctx.num_groups = num_groups
input_list = [
torch.zeros_like(input) for k in range(du.get_local_size())
]
dist.all_gather(
input_list, input, async_op=False, group=du._LOCAL_PROCESS_GROUP
)
inputs = torch.stack(input_list, dim=0)
if num_groups > 1:
rank = du.get_local_rank()
group_idx = rank // num_sync_devices
inputs = inputs[
group_idx
* num_sync_devices : (group_idx + 1)
* num_sync_devices
]
inputs = torch.sum(inputs, dim=0)
return inputs
def backward(ctx, grad_output):
"""
Perform backwarding, gathering the gradients across different process/ GPU
group.
"""
grad_output_list = [
torch.zeros_like(grad_output) for k in range(du.get_local_size())
]
dist.all_gather(
grad_output_list,
grad_output,
async_op=False,
group=du._LOCAL_PROCESS_GROUP,
)
grads = torch.stack(grad_output_list, dim=0)
if ctx.num_groups > 1:
rank = du.get_local_rank()
group_idx = rank // ctx.num_sync_devices
grads = grads[
group_idx
* ctx.num_sync_devices : (group_idx + 1)
* ctx.num_sync_devices
]
grads = torch.sum(grads, dim=0)
return grads, None, None
def _batch_unshuffle(self, x, idx_restore):
if self.num_gpus > 1:
if self.cfg.CONTRASTIVE.LOCAL_SHUFFLE_BN:
x = du.cat_all_gather(x, local=True)
gpu_idx = du.get_local_rank()
else:
x = du.cat_all_gather(x)
gpu_idx = torch.distributed.get_rank()
else:
gpu_idx = 0
idx = idx_restore[gpu_idx, :]
x = x[idx]
return x
def _simclr_precompute_pos_neg_mask_multi(self):
# computed once at the beginning of training
distributed = self.cfg.CONTRASTIVE.SIMCLR_DIST_ON
if distributed:
total_images = self.cfg.TRAIN.BATCH_SIZE * self.cfg.NUM_SHARDS
world_size = du.get_world_size()
rank = du.get_rank()
else:
total_images = self.cfg.TRAIN.BATCH_SIZE
world_size = du.get_local_size()
rank = du.get_local_rank()
local_orig_images = total_images // world_size
local_crops = local_orig_images * self.num_crops
pos_temps = []
for d in np.arange(self.num_crops):
pos_temp, neg_temp = [], []
for i in range(world_size):
if i == rank:
pos = np.eye(local_crops,
k=d * local_orig_images) + np.eye(
local_crops,
k=-local_crops + d * local_orig_images)
neg = np.ones((local_crops, local_crops))
else:
pos = np.zeros((local_crops, local_crops))
neg = np.zeros((local_crops, local_crops))
pos_temp.append(pos)
neg_temp.append(neg)
pos_temps.append(np.hstack(pos_temp))
neg_temp = np.hstack(neg_temp)
pos_mask = []
for i in range(self.num_crops - 1):
pos_mask.append(torch.from_numpy(pos_temps[1 + i]))
neg_mask = torch.from_numpy(neg_temp - sum(pos_temps))
if self.num_gpus:
for i in range(len(pos_mask)):
pos_mask[i] = pos_mask[i].cuda(non_blocking=True)
neg_mask = neg_mask.cuda(non_blocking=True)
self.pos_mask, self.neg_mask = pos_mask, neg_mask
def _batch_shuffle(self, x):
if len(x) == 2:
another_crop = True
else:
another_crop = False
if another_crop:
x, x_crop = x[0], x[1]
else:
x = x[0]
world_size = self.cfg.NUM_GPUS * self.cfg.NUM_SHARDS
if self.num_gpus > 1:
if self.cfg.CONTRASTIVE.LOCAL_SHUFFLE_BN:
x = du.cat_all_gather(x, local=True)
if another_crop:
x_crop = du.cat_all_gather(x_crop, local=True)
world_size = du.get_local_size()
gpu_idx = du.get_local_rank()
else:
x = du.cat_all_gather(x)
if another_crop:
x_crop = du.cat_all_gather(x_crop)
gpu_idx = torch.distributed.get_rank()
idx_randperm = torch.randperm(x.shape[0]).cuda()
if self.num_gpus > 1:
torch.distributed.broadcast(idx_randperm, src=0)
else:
gpu_idx = 0
idx_randperm = idx_randperm.view(world_size, -1)
x = x[idx_randperm[gpu_idx, :]]
if another_crop:
x_crop = x_crop[idx_randperm[gpu_idx, :]]
idx_restore = torch.argsort(idx_randperm.view(-1))
idx_restore = idx_restore.view(world_size, -1)
if another_crop:
return [x, x_crop], idx_restore
else:
return [x], idx_restore
def train_epoch(self,
train_loader,
model,
optimizer,
train_meter,
cur_epoch,
cfg,
writer=None):
"""
Perform the video training for one epoch.
Args:
train_loader (loader): video training loader.
model (model): the video model to train.
optimizer (optim): the optimizer to perform optimization on the model's
parameters.
train_meter (TrainMeter): training meters to log the training performance.
cur_epoch (int): current epoch of training.
cfg (CfgNode): configs. Details can be found in
slowfast/config/defaults.py
writer (TensorboardWriter, optional): TensorboardWriter object
to writer Tensorboard log.
"""
# Enable train mode.
model.train()
train_meter.iter_tic()
data_size = len(train_loader)
start = time.time()
btch = cfg.TRAIN.BATCH_SIZE * self.cfg.NUM_SHARDS
rankE = os.environ.get("RANK", None)
worldE = os.environ.get("WORLD_SIZE", None)
dSize = data_size * btch
self.logger.info(
"Train Epoch {} dLen {} Batch {} dSize {} localRank {} rank {} {} world {} {}"
.format(cur_epoch, data_size, btch, dSize, du.get_local_rank(),
du.get_rank(), rankE, du.get_world_size(), worldE))
tot = 0
first = True
predsAll = []
labelsAll = []
for cur_iter, (inputs, labels, _, meta) in enumerate(train_loader):
# Transfer the data to the current GPU device.
tot += len(labels)
if isinstance(inputs, (list, )):
if first:
self.logger.info(
"rank {} LEN {} {} shape Slow {} Fast {} {} tot {}".
format(du.get_rank(), len(labels), len(inputs),
inputs[0].shape, inputs[1].shape,
labels[0].shape, tot))
first = False
for i in range(len(inputs)):
inputs[i] = inputs[i].cuda(non_blocking=True)
else:
if first:
self.logger.info(
"rank {} LEN {} shape {} {} tot {}".format(
du.get_rank(), len(labels), inputs.shape,
labels[0].shape, tot))
first = False
inputs = inputs.cuda(non_blocking=True)
labels = labels.cuda()
for key, val in meta.items():
if isinstance(val, (list, )):
for i in range(len(val)):
val[i] = val[i].cuda(non_blocking=True)
else:
meta[key] = val.cuda(non_blocking=True)
# Update the learning rate.
lr = optim.get_epoch_lr(cur_epoch + float(cur_iter) / data_size,
cfg)
optim.set_lr(optimizer, lr)
if cfg.DETECTION.ENABLE:
# Compute the predictions.
preds = model(inputs, meta["boxes"])
else:
# Perform the forward pass.
preds = model(inputs)
# Explicitly declare reduction to mean.
loss_fun = losses.get_loss_func(
cfg.MODEL.LOSS_FUNC)(reduction="mean")
# Compute the loss.
loss = loss_fun(preds, labels)
# check Nan Loss.
misc.check_nan_losses(loss)
# Perform the backward pass.
optimizer.zero_grad()
loss.backward()
# Update the parameters.
optimizer.step()
if cfg.DETECTION.ENABLE:
if cfg.NUM_GPUS > 1:
loss = du.all_reduce([loss])[0]
loss = loss.item()
train_meter.iter_toc()
# Update and log stats.
train_meter.update_stats(None, None, None, loss, lr)
# write to tensorboard format if available.
if writer is not None:
writer.add_scalars(
{
"Train/loss": loss,
"Train/lr": lr
},
global_step=data_size * cur_epoch + cur_iter,
)
ite = data_size * cur_epoch + cur_iter
if du.is_master_proc():
self.logger.log_row(name='TrainLoss',
iter=ite,
loss=loss,
description="train loss")
self.logger.log_row(name='TrainLr',
iter=ite,
lr=lr,
description="train learn rate")
else:
top1_err, top5_err = None, None
if cfg.DATA.MULTI_LABEL:
# Gather all the predictions across all the devices.
if cfg.NUM_GPUS > 1:
[loss] = du.all_reduce([loss])
loss = loss.item()
else:
# Binary classifier - save preds / labels for metrics
if cfg.MODEL.NUM_CLASSES == 2:
predsAll.extend(preds.detach().cpu().numpy()[:, -1])
labelsAll.extend(labels.detach().cpu().numpy())
# Compute the errors.
num_topks_correct = metrics.topks_correct(
preds, labels, (1, min(5, cfg.MODEL.NUM_CLASSES)))
top1_err, top5_err = [(1.0 - x / preds.size(0)) * 100.0
for x in num_topks_correct]
# Gather all the predictions across all the devices.
if cfg.NUM_GPUS > 1:
loss, top1_err, top5_err = du.all_reduce(
[loss, top1_err, top5_err])
# Copy the stats from GPU to CPU (sync point).
loss, top1_err, top5_err = (
loss.item(),
top1_err.item(),
top5_err.item(),
)
train_meter.iter_toc()
# Update and log stats.
# self.logger.info("UPDATING stat {} {} {}".format(inputs[0].size(0), cfg.NUM_GPUS, inputs[0].size(0) * cfg.NUM_GPUS))
train_meter.update_stats(top1_err, top5_err, loss, lr,
inputs[0].size(0) * cfg.NUM_GPUS)
# write to tensorboard format if available.
if writer is not None:
writer.add_scalars(
{
"Train/loss": loss,
"Train/lr": lr,
"Train/Top1_err": top1_err,
"Train/Top5_err": top5_err,
},
global_step=data_size * cur_epoch + cur_iter,
)
stats = train_meter.log_iter_stats(cur_epoch, cur_iter, predsAll,
labelsAll)
ite = dSize * cur_epoch + btch * (cur_iter + 1)
self.plotStats(stats, ite, 'TrainIter')
train_meter.iter_tic()
if du.is_master_proc() and cfg.LOG_MODEL_INFO:
misc.log_model_info(model, cfg, use_train_input=True)
# Log epoch stats.
gathered = du.all_gather([
torch.tensor(predsAll).to(torch.device("cuda")),
torch.tensor(labelsAll).to(torch.device("cuda"))
])
stats = train_meter.log_epoch_stats(cur_epoch,
gathered[0].detach().cpu().numpy(),
gathered[1].detach().cpu().numpy())
ite = (cur_epoch + 1) * dSize
self.plotStats(stats, ite, 'TrainEpoch')
train_meter.reset()
end = time.time()
el = end - start
totAll = du.all_reduce([torch.tensor(tot).cuda()], average=False)
tSum = totAll[0].item()
elT = torch.tensor(el).cuda()
elMax = du.all_reduce([elT], op=dist.ReduceOp.MAX,
average=False)[0].item()
jobRate = tSum / elMax
self.logger.info(
"totSampCnt {} workerSampCnt {} eTimeMax {} eTimeWorker {} SampPerSecJob {:.1f} SampPerSecWorker {:.1f}"
.format(tSum, tot, elMax, el, jobRate, tot / el))
return jobRate
def eval_epoch(self,
val_loader,
model,
val_meter,
cur_epoch,
cfg,
writer=None):
"""
Evaluate the model on the val set.
Args:
val_loader (loader): data loader to provide validation data.
model (model): model to evaluate the performance.
val_meter (ValMeter): meter instance to record and calculate the metrics.
cur_epoch (int): number of the current epoch of training.
cfg (CfgNode): configs. Details can be found in
slowfast/config/defaults.py
writer (TensorboardWriter, optional): TensorboardWriter object
to writer Tensorboard log.
"""
# Evaluation mode enabled. The running stats would not be updated.
model.eval()
data_size = len(val_loader)
btch = cfg.TRAIN.BATCH_SIZE * self.cfg.NUM_SHARDS
rankE = os.environ.get("RANK", None)
worldE = os.environ.get("WORLD_SIZE", None)
dSize = data_size * btch
self.logger.info(
"Val Epoch {} dLen {} Batch {} dSize {} localRank {} rank {} {} world {} {}"
.format(cur_epoch, data_size, btch, dSize, du.get_local_rank(),
du.get_rank(), rankE, du.get_world_size(), worldE))
val_meter.iter_tic()
predsAll = []
labelsAll = []
data_size = len(val_loader)
for cur_iter, (inputs, labels, _, meta) in enumerate(val_loader):
# Transferthe data to the current GPU device.
if isinstance(inputs, (list, )):
for i in range(len(inputs)):
inputs[i] = inputs[i].cuda(non_blocking=True)
else:
inputs = inputs.cuda(non_blocking=True)
labels = labels.cuda()
for key, val in meta.items():
if isinstance(val, (list, )):
for i in range(len(val)):
val[i] = val[i].cuda(non_blocking=True)
else:
meta[key] = val.cuda(non_blocking=True)
if cfg.DETECTION.ENABLE:
# Compute the predictions.
preds = model(inputs, meta["boxes"])
preds = preds.cpu()
ori_boxes = meta["ori_boxes"].cpu()
metadata = meta["metadata"].cpu()
if cfg.NUM_GPUS > 1:
preds = torch.cat(du.all_gather_unaligned(preds), dim=0)
ori_boxes = torch.cat(du.all_gather_unaligned(ori_boxes),
dim=0)
metadata = torch.cat(du.all_gather_unaligned(metadata),
dim=0)
val_meter.iter_toc()
# Update and log stats.
val_meter.update_stats(preds.cpu(), ori_boxes.cpu(),
metadata.cpu())
else:
preds = model(inputs)
if cfg.DATA.MULTI_LABEL:
if cfg.NUM_GPUS > 1:
preds, labels = du.all_gather([preds, labels])
else:
if cfg.MODEL.NUM_CLASSES == 2:
predsAll.extend(preds.detach().cpu().numpy()[:, -1])
labelsAll.extend(labels.detach().cpu().numpy())
# Compute the errors.
num_topks_correct = metrics.topks_correct(
preds, labels, (1, min(5, cfg.MODEL.NUM_CLASSES)))
# Combine the errors across the GPUs.
top1_err, top5_err = [(1.0 - x / preds.size(0)) * 100.0
for x in num_topks_correct]
if cfg.NUM_GPUS > 1:
top1_err, top5_err = du.all_reduce(
[top1_err, top5_err])
# Copy the errors from GPU to CPU (sync point).
top1_err, top5_err = top1_err.item(), top5_err.item()
val_meter.iter_toc()
# Update and log stats.
val_meter.update_stats(top1_err, top5_err,
inputs[0].size(0) * cfg.NUM_GPUS)
# write to tensorboard format if available.
if writer is not None:
writer.add_scalars(
{
"Val/Top1_err": top1_err,
"Val/Top5_err": top5_err
},
global_step=len(val_loader) * cur_epoch + cur_iter,
)
if du.is_master_proc():
ite = len(val_loader) * cur_epoch + cur_iter
self.logger.log_row(name='ValTop1',
iter=ite,
lr=top1_err,
description="Top 1 Err")
self.logger.log_row(name='ValTop5',
iter=ite,
lr=top5_err,
description="Top 5 Err")
val_meter.update_predictions(preds, labels)
stats = val_meter.log_iter_stats(cur_epoch, cur_iter, predsAll,
labelsAll)
ite = dSize * cur_epoch + btch * (cur_iter + 1)
self.plotStats(stats, ite, 'ValIter')
val_meter.iter_tic()
# Log epoch stats.
gathered = du.all_gather([
torch.tensor(predsAll).to(torch.device("cuda")),
torch.tensor(labelsAll).to(torch.device("cuda"))
])
stats = val_meter.log_epoch_stats(cur_epoch,
gathered[0].detach().cpu().numpy(),
gathered[1].detach().cpu().numpy())
ite = (cur_epoch + 1) * dSize
self.plotStats(stats, ite, 'ValEpoch')
# write to tensorboard format if available.
if writer is not None:
if cfg.DETECTION.ENABLE:
writer.add_scalars({"Val/mAP": val_meter.full_map},
global_step=cur_epoch)
all_preds_cpu = [
pred.clone().detach().cpu() for pred in val_meter.all_preds
]
all_labels_cpu = [
label.clone().detach().cpu() for label in val_meter.all_labels
]
# plotScatter(all_preds_cpu, all_labels_cpu, "Epoch_{}".format(cur_epoch))
# writer.plot_eval(
# preds=all_preds_cpu, labels=all_labels_cpu, global_step=cur_epoch
# )
val_meter.reset()