def test(cfg):
"""
Perform multi-view testing on the pretrained video model.
Args:
cfg (CfgNode): configs. Details can be found in
slowfast/config/defaults.py
"""
# Set up environment.
du.init_distributed_training(cfg)
# Set random seed from configs.
np.random.seed(cfg.RNG_SEED)
torch.manual_seed(cfg.RNG_SEED)
# Setup logging format.
logging.setup_logging(cfg.OUTPUT_DIR)
# Print config.
logger.info("Test with config:")
logger.info(cfg)
# Build the video model and print model statistics.
model = build_model(cfg)
if du.is_master_proc() and cfg.LOG_MODEL_INFO:
misc.log_model_info(model, cfg, use_train_input=False)
cu.load_test_checkpoint(cfg, model)
# Create video testing loaders.
test_loader = loader.construct_loader(cfg, "test")
logger.info("Testing model for {} iterations".format(len(test_loader)))
if cfg.DETECTION.ENABLE:
assert cfg.NUM_GPUS == cfg.TEST.BATCH_SIZE or cfg.NUM_GPUS == 0
test_meter = AVAMeter(len(test_loader), cfg, mode="test")
else:
assert (
test_loader.dataset.num_videos %
(cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS) == 0)
# Create meters for multi-view testing.
test_meter = TestMeter(
test_loader.dataset.num_videos //
(cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS),
cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS,
cfg.MODEL.NUM_CLASSES,
len(test_loader),
cfg.DATA.MULTI_LABEL,
cfg.DATA.ENSEMBLE_METHOD,
)
# Set up writer for logging to Tensorboard format.
if cfg.TENSORBOARD.ENABLE and du.is_master_proc(
cfg.NUM_GPUS * cfg.NUM_SHARDS):
writer = tb.TensorboardWriter(cfg)
else:
writer = None
# # Perform multi-view test on the entire dataset.
test_meter = perform_test(test_loader, model, test_meter, cfg, writer)
if writer is not None:
writer.close()
def visualize(cfg):
"""
Perform layer weights and activations visualization on the model.
Args:
cfg (CfgNode): configs. Details can be found in
slowfast/config/defaults.py
"""
if cfg.TENSORBOARD.ENABLE and cfg.TENSORBOARD.MODEL_VIS.ENABLE:
# Set up environment.
du.init_distributed_training(cfg)
# Set random seed from configs.
np.random.seed(cfg.RNG_SEED)
torch.manual_seed(cfg.RNG_SEED)
# Setup logging format.
logging.setup_logging(cfg.OUTPUT_DIR)
# Print config.
logger.info("Model Visualization with config:")
logger.info(cfg)
# Build the video model and print model statistics.
model = build_model(cfg)
if du.is_master_proc() and cfg.LOG_MODEL_INFO:
misc.log_model_info(model, cfg, is_train=False)
cu.load_test_checkpoint(cfg, model)
# Create video testing loaders.
vis_loader = loader.construct_loader(cfg, "test")
logger.info(
"Visualize model for {} data points".format(len(vis_loader))
)
if cfg.DETECTION.ENABLE:
assert cfg.NUM_GPUS == cfg.TEST.BATCH_SIZE
# Set up writer for logging to Tensorboard format.
if du.is_master_proc(cfg.NUM_GPUS * cfg.NUM_SHARDS):
writer = tb.TensorboardWriter(cfg)
else:
writer = None
# Run visualization on the model
run_visualization(vis_loader, model, cfg, writer)
if writer is not None:
writer.close()
def test(cfg):
"""
Perform multi-view testing on the pretrained video model.
Args:
cfg (CfgNode): configs. Details can be found in
slowfast/config/defaults.py
"""
# Set up environment.
du.init_distributed_training(cfg)
# Set random seed from configs.
np.random.seed(cfg.RNG_SEED)
torch.manual_seed(cfg.RNG_SEED)
# Setup logging format.
logging.setup_logging(cfg.OUTPUT_DIR)
# Print config.
logger.info("Test with config:")
logger.info(cfg)
# Build the video model and print model statistics.
model = build_model(cfg)
cu.load_test_checkpoint(cfg, model)
# Create video testing loaders.
test_loader = loader.construct_loader(cfg, "test")
logger.info("Testing model for {} iterations".format(len(test_loader)))
# Create meters for loss tracking
test_meter = TrainMeter(test_loader.dataset.num_videos, cfg)
# Set up writer for logging to Tensorboard format.
if cfg.TENSORBOARD.ENABLE and du.is_master_proc(
cfg.NUM_GPUS * cfg.NUM_SHARDS
):
writer = tb.TensorboardWriter(cfg)
else:
writer = None
# # Perform multi-view test on the entire dataset.
test_meter = perform_test(test_loader, model, test_meter, cfg, writer)
if writer is not None:
writer.close()
def __init__(self, cfg):
"""
Args:
cfg (CfgNode): configs. Details can be found in
slowfast/config/defaults.py
"""
self.cfg = cfg
self.class_names, _, self.subset = get_class_names(
cfg.TENSORBOARD.CLASS_NAMES_PATH,
subset_path=cfg.TENSORBOARD.WRONG_PRED_VIS.SUBSET_PATH,
)
if self.subset is not None:
self.subset = set(self.subset)
self.num_class = cfg.MODEL.NUM_CLASSES
self.video_vis = VideoVisualizer(
cfg.MODEL.NUM_CLASSES,
cfg.TENSORBOARD.CLASS_NAMES_PATH,
1,
cfg.TENSORBOARD.MODEL_VIS.COLORMAP,
)
self.tag = cfg.TENSORBOARD.WRONG_PRED_VIS.TAG
self.writer = tb.TensorboardWriter(cfg)
self.model_incorrect_classes = set()
def train(cfg):
"""
Train a video model for many epochs on train set and evaluate it on val set.
Args:
cfg (CfgNode): configs. Details can be found in
slowfast/config/defaults.py
"""
# Set up environment.
du.init_distributed_training(cfg)
# Set random seed from configs.
np.random.seed(cfg.RNG_SEED)
torch.manual_seed(cfg.RNG_SEED)
# Setup logging format.
logging.setup_logging(cfg.OUTPUT_DIR)
# Init multigrid.
multigrid = None
if cfg.MULTIGRID.LONG_CYCLE or cfg.MULTIGRID.SHORT_CYCLE:
multigrid = MultigridSchedule()
cfg = multigrid.init_multigrid(cfg)
if cfg.MULTIGRID.LONG_CYCLE:
cfg, _ = multigrid.update_long_cycle(cfg, cur_epoch=0)
# Print config.
logger.info("Train with config:")
logger.info(pprint.pformat(cfg))
# Build the video model and print model statistics.
model = build_model(cfg)
if du.is_master_proc() and cfg.LOG_MODEL_INFO:
misc.log_model_info(model, cfg, use_train_input=True)
# Construct the optimizer.
optimizer = optim.construct_optimizer(model, cfg)
# Load a checkpoint to resume training if applicable.
start_epoch = cu.load_train_checkpoint(cfg, model, optimizer)
# Create the video train and val loaders.
train_loader = loader.construct_loader(cfg, "train")
val_loader = loader.construct_loader(cfg, "val")
precise_bn_loader = (loader.construct_loader(
cfg, "train", is_precise_bn=True)
if cfg.BN.USE_PRECISE_STATS else None)
# Create meters.
if cfg.DETECTION.ENABLE:
train_meter = AVAMeter(len(train_loader), cfg, mode="train")
val_meter = AVAMeter(len(val_loader), cfg, mode="val")
else:
train_meter = TrainMeter(len(train_loader), cfg)
val_meter = ValMeter(len(val_loader), cfg)
# set up writer for logging to Tensorboard format.
if cfg.TENSORBOARD.ENABLE and du.is_master_proc(
cfg.NUM_GPUS * cfg.NUM_SHARDS):
writer = tb.TensorboardWriter(cfg)
else:
writer = None
# Perform the training loop.
logger.info("Start epoch: {}".format(start_epoch + 1))
for cur_epoch in range(start_epoch, cfg.SOLVER.MAX_EPOCH):
if cfg.MULTIGRID.LONG_CYCLE:
cfg, changed = multigrid.update_long_cycle(cfg, cur_epoch)
if changed:
(
model,
optimizer,
train_loader,
val_loader,
precise_bn_loader,
train_meter,
val_meter,
) = build_trainer(cfg)
# Load checkpoint.
if cu.has_checkpoint(cfg.OUTPUT_DIR):
last_checkpoint = cu.get_last_checkpoint(cfg.OUTPUT_DIR)
assert "{:05d}.pyth".format(cur_epoch) in last_checkpoint
else:
last_checkpoint = cfg.TRAIN.CHECKPOINT_FILE_PATH
logger.info("Load from {}".format(last_checkpoint))
cu.load_checkpoint(last_checkpoint, model, cfg.NUM_GPUS > 1,
optimizer)
# Shuffle the dataset.
loader.shuffle_dataset(train_loader, cur_epoch)
# Train for one epoch.
train_epoch(train_loader, model, optimizer, train_meter, cur_epoch,
cfg, writer)
is_checkp_epoch = (cu.is_checkpoint_epoch(
cfg,
cur_epoch,
None if multigrid is None else multigrid.schedule,
))
is_eval_epoch = misc.is_eval_epoch(
cfg, cur_epoch, None if multigrid is None else multigrid.schedule)
# Compute precise BN stats.
if ((is_checkp_epoch or is_eval_epoch) and cfg.BN.USE_PRECISE_STATS
and len(get_bn_modules(model)) > 0):
calculate_and_update_precise_bn(
precise_bn_loader,
model,
min(cfg.BN.NUM_BATCHES_PRECISE, len(precise_bn_loader)),
cfg.NUM_GPUS > 0,
)
_ = misc.aggregate_sub_bn_stats(model)
# Save a checkpoint.
if is_checkp_epoch:
cu.save_checkpoint(cfg.OUTPUT_DIR, model, optimizer, cur_epoch,
cfg)
# Evaluate the model on validation set.
if is_eval_epoch:
eval_epoch(val_loader, model, val_meter, cur_epoch, cfg, writer)
if writer is not None:
writer.close()
def visualize(cfg):
"""
Perform layer weights and activations visualization on the model.
Args:
cfg (CfgNode): configs. Details can be found in
slowfast/config/defaults.py
"""
if cfg.TENSORBOARD.ENABLE and (cfg.TENSORBOARD.MODEL_VIS.ENABLE
or cfg.TENSORBOARD.WRONG_PRED_VIS.ENABLE):
# Set up environment.
du.init_distributed_training(cfg)
# Set random seed from configs.
np.random.seed(cfg.RNG_SEED)
torch.manual_seed(cfg.RNG_SEED)
# Setup logging format.
logging.setup_logging(cfg.OUTPUT_DIR)
# Print config.
logger.info("Model Visualization with config:")
logger.info(cfg)
# Build the video model and print model statistics.
model = build_model(cfg)
model.eval()
if du.is_master_proc() and cfg.LOG_MODEL_INFO:
misc.log_model_info(model, cfg, use_train_input=False)
cu.load_test_checkpoint(cfg, model)
# Create video testing loaders.
vis_loader = loader.construct_loader(cfg, "test")
if cfg.DETECTION.ENABLE:
assert cfg.NUM_GPUS == cfg.TEST.BATCH_SIZE or cfg.NUM_GPUS == 0
# Set up writer for logging to Tensorboard format.
if du.is_master_proc(cfg.NUM_GPUS * cfg.NUM_SHARDS):
writer = tb.TensorboardWriter(cfg)
else:
writer = None
if cfg.TENSORBOARD.PREDICTIONS_PATH != "":
assert not cfg.DETECTION.ENABLE, "Detection is not supported."
logger.info(
"Visualizing class-level performance from saved results...")
if writer is not None:
with g_pathmgr.open(cfg.TENSORBOARD.PREDICTIONS_PATH,
"rb") as f:
preds, labels = pickle.load(f, encoding="latin1")
writer.plot_eval(preds, labels)
if cfg.TENSORBOARD.MODEL_VIS.ENABLE:
if cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.ENABLE:
assert (
not cfg.DETECTION.ENABLE
), "Detection task is currently not supported for Grad-CAM visualization."
if cfg.MODEL.ARCH in cfg.MODEL.SINGLE_PATHWAY_ARCH:
assert (
len(cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST) == 1
), "The number of chosen CNN layers must be equal to the number of pathway(s), given {} layer(s).".format(
len(cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST))
elif cfg.MODEL.ARCH in cfg.MODEL.MULTI_PATHWAY_ARCH:
assert (
len(cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST) == 2
), "The number of chosen CNN layers must be equal to the number of pathway(s), given {} layer(s).".format(
len(cfg.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST))
else:
raise NotImplementedError(
"Model arch {} is not in {}".format(
cfg.MODEL.ARCH,
cfg.MODEL.SINGLE_PATHWAY_ARCH +
cfg.MODEL.MULTI_PATHWAY_ARCH,
))
logger.info("Visualize model analysis for {} iterations".format(
len(vis_loader)))
# Run visualization on the model
run_visualization(vis_loader, model, cfg, writer)
if cfg.TENSORBOARD.WRONG_PRED_VIS.ENABLE:
logger.info("Visualize Wrong Predictions for {} iterations".format(
len(vis_loader)))
perform_wrong_prediction_vis(vis_loader, model, cfg)
if writer is not None:
writer.close()
def test(cfg):
"""
Perform multi-view testing on the pretrained video model.
Args:
cfg (CfgNode): configs. Details can be found in
slowfast/config/defaults.py
"""
# Set up environment.
du.init_distributed_training(cfg)
# Set random seed from configs.
np.random.seed(cfg.RNG_SEED)
torch.manual_seed(cfg.RNG_SEED)
# Setup logging format.
logging.setup_logging(cfg.OUTPUT_DIR)
# Print config.
logger.info("Test with config:")
logger.info(cfg)
# Build the video model and print model statistics.
model = build_model(cfg)
out_str_prefix = "lin" if cfg.MODEL.DETACH_FINAL_FC else ""
if du.is_master_proc() and cfg.LOG_MODEL_INFO:
misc.log_model_info(model, cfg, use_train_input=False)
if (cfg.TASK == "ssl" and cfg.MODEL.MODEL_NAME == "ContrastiveModel"
and cfg.CONTRASTIVE.KNN_ON):
train_loader = loader.construct_loader(cfg, "train")
out_str_prefix = "knn"
if hasattr(model, "module"):
model.module.init_knn_labels(train_loader)
else:
model.init_knn_labels(train_loader)
cu.load_test_checkpoint(cfg, model)
# Create video testing loaders.
test_loader = loader.construct_loader(cfg, "test")
logger.info("Testing model for {} iterations".format(len(test_loader)))
if cfg.DETECTION.ENABLE:
assert cfg.NUM_GPUS == cfg.TEST.BATCH_SIZE or cfg.NUM_GPUS == 0
test_meter = AVAMeter(len(test_loader), cfg, mode="test")
else:
assert (
test_loader.dataset.num_videos %
(cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS) == 0)
# Create meters for multi-view testing.
test_meter = TestMeter(
test_loader.dataset.num_videos //
(cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS),
cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS,
cfg.MODEL.NUM_CLASSES if not cfg.TASK == "ssl" else
cfg.CONTRASTIVE.NUM_CLASSES_DOWNSTREAM,
len(test_loader),
cfg.DATA.MULTI_LABEL,
cfg.DATA.ENSEMBLE_METHOD,
)
# Set up writer for logging to Tensorboard format.
if cfg.TENSORBOARD.ENABLE and du.is_master_proc(
cfg.NUM_GPUS * cfg.NUM_SHARDS):
writer = tb.TensorboardWriter(cfg)
else:
writer = None
# # Perform multi-view test on the entire dataset.
test_meter = perform_test(test_loader, model, test_meter, cfg, writer)
if writer is not None:
writer.close()
result_string = (
"_a{}{}{} Top1 Acc: {} Top5 Acc: {} MEM: {:.2f} dataset: {}{}"
"".format(
out_str_prefix,
cfg.TEST.DATASET[0],
test_meter.stats["top1_acc"],
test_meter.stats["top1_acc"],
test_meter.stats["top5_acc"],
misc.gpu_mem_usage(),
cfg.TEST.DATASET[0],
cfg.MODEL.NUM_CLASSES,
))
logger.info("testing done: {}".format(result_string))
return result_string
def test(cfg):
"""
Perform multi-view testing on the pretrained audio model.
Args:
cfg (CfgNode): configs. Details can be found in
slowfast/config/defaults.py
"""
# Set up environment.
du.init_distributed_training(cfg)
# Set random seed from configs.
np.random.seed(cfg.RNG_SEED)
torch.manual_seed(cfg.RNG_SEED)
# Setup logging format.
logging.setup_logging(cfg.OUTPUT_DIR)
# Print config.
logger.info("Test with config:")
logger.info(cfg)
# Build the audio model and print model statistics.
model = build_model(cfg)
if du.is_master_proc() and cfg.LOG_MODEL_INFO:
misc.log_model_info(model, cfg)
cu.load_test_checkpoint(cfg, model)
# Create audio testing loaders.
test_loader = loader.construct_loader(cfg, "test")
logger.info("Testing model for {} iterations".format(len(test_loader)))
assert (
len(test_loader.dataset)
% cfg.TEST.NUM_ENSEMBLE_VIEWS
== 0
)
# Create meters for multi-view testing.
if cfg.TEST.DATASET == 'epickitchens':
test_meter = EPICTestMeter(
len(test_loader.dataset)
// cfg.TEST.NUM_ENSEMBLE_VIEWS,
cfg.TEST.NUM_ENSEMBLE_VIEWS,
cfg.MODEL.NUM_CLASSES,
len(test_loader),
cfg.DATA.ENSEMBLE_METHOD,
)
else:
test_meter = TestMeter(
len(test_loader.dataset)
// cfg.TEST.NUM_ENSEMBLE_VIEWS,
cfg.TEST.NUM_ENSEMBLE_VIEWS,
cfg.MODEL.NUM_CLASSES[0],
len(test_loader),
cfg.DATA.MULTI_LABEL,
cfg.DATA.ENSEMBLE_METHOD,
)
# Set up writer for logging to Tensorboard format.
if cfg.TENSORBOARD.ENABLE and du.is_master_proc(
cfg.NUM_GPUS * cfg.NUM_SHARDS
):
writer = tb.TensorboardWriter(cfg)
else:
writer = None
# # Perform multi-view test on the entire dataset.
test_meter, preds, preds_clips, labels, metadata = perform_test(test_loader, model, test_meter, cfg, writer)
if du.is_master_proc():
if cfg.TEST.DATASET == 'epickitchens':
results = {'verb_output': preds[0],
'noun_output': preds[1],
'narration_id': metadata}
scores_path = os.path.join(cfg.OUTPUT_DIR, 'scores')
if not os.path.exists(scores_path):
os.makedirs(scores_path)
file_path = os.path.join(scores_path, cfg.EPICKITCHENS.TEST_SPLIT+'.pkl')
pickle.dump(results, open(file_path, 'wb'))
else:
if cfg.TEST.DATASET == 'vggsound':
get_stats(preds, labels)
results = {'scores': preds, 'labels': labels}
scores_path = os.path.join(cfg.OUTPUT_DIR, 'scores')
if not os.path.exists(scores_path):
os.makedirs(scores_path)
file_path = os.path.join(scores_path, 'test.pkl')
pickle.dump(results, open(file_path, 'wb'))
if writer is not None:
writer.close()
def train(cfg):
"""
Train an audio model for many epochs on train set and evaluate it on val set.
Args:
cfg (CfgNode): configs. Details can be found in
slowfast/config/defaults.py
"""
# Set up environment.
du.init_distributed_training(cfg)
# Set random seed from configs.
np.random.seed(cfg.RNG_SEED)
torch.manual_seed(cfg.RNG_SEED)
# Setup logging format.
logging.setup_logging(cfg.OUTPUT_DIR)
# Print config.
logger.info("Train with config:")
logger.info(pprint.pformat(cfg))
# Build the audio model and print model statistics.
model = build_model(cfg)
if du.is_master_proc() and cfg.LOG_MODEL_INFO:
misc.log_model_info(model, cfg)
if cfg.BN.FREEZE:
model.module.freeze_fn(
'bn_parameters') if cfg.NUM_GPUS > 1 else model.freeze_fn(
'bn_parameters')
# Construct the optimizer.
optimizer = optim.construct_optimizer(model, cfg)
# Load a checkpoint to resume training if applicable.
start_epoch = cu.load_train_checkpoint(cfg, model, optimizer)
# Create the audio train and val loaders.
if cfg.TRAIN.DATASET != 'epickitchens' or not cfg.EPICKITCHENS.TRAIN_PLUS_VAL:
train_loader = loader.construct_loader(cfg, "train")
val_loader = loader.construct_loader(cfg, "val")
precise_bn_loader = (loader.construct_loader(cfg, "train")
if cfg.BN.USE_PRECISE_STATS else None)
else:
train_loader = loader.construct_loader(cfg, "train+val")
val_loader = loader.construct_loader(cfg, "val")
precise_bn_loader = (loader.construct_loader(cfg, "train+val")
if cfg.BN.USE_PRECISE_STATS else None)
# Create meters.
if cfg.TRAIN.DATASET == 'epickitchens':
train_meter = EPICTrainMeter(len(train_loader), cfg)
val_meter = EPICValMeter(len(val_loader), cfg)
else:
train_meter = TrainMeter(len(train_loader), cfg)
val_meter = ValMeter(len(val_loader), cfg)
# set up writer for logging to Tensorboard format.
if cfg.TENSORBOARD.ENABLE and du.is_master_proc(
cfg.NUM_GPUS * cfg.NUM_SHARDS):
writer = tb.TensorboardWriter(cfg)
else:
writer = None
if cfg.WANDB.ENABLE and du.is_master_proc(cfg.NUM_GPUS * cfg.NUM_SHARDS):
wandb_log = True
if cfg.TRAIN.AUTO_RESUME and cfg.WANDB.RUN_ID != "":
wandb.init(project='slowfast',
config=cfg,
sync_tensorboard=True,
resume=cfg.WANDB.RUN_ID)
else:
wandb.init(project='slowfast', config=cfg, sync_tensorboard=True)
wandb.watch(model)
else:
wandb_log = False
# Perform the training loop.
logger.info("Start epoch: {}".format(start_epoch + 1))
for cur_epoch in range(start_epoch, cfg.SOLVER.MAX_EPOCH):
# Shuffle the dataset.
loader.shuffle_dataset(train_loader, cur_epoch)
# Train for one epoch.
train_epoch(train_loader, model, optimizer, train_meter, cur_epoch,
cfg, writer, wandb_log)
is_checkp_epoch = cu.is_checkpoint_epoch(
cfg,
cur_epoch,
)
is_eval_epoch = misc.is_eval_epoch(
cfg,
cur_epoch,
)
# Compute precise BN stats.
if ((is_checkp_epoch or is_eval_epoch) and cfg.BN.USE_PRECISE_STATS
and len(get_bn_modules(model)) > 0):
calculate_and_update_precise_bn(
precise_bn_loader,
model,
min(cfg.BN.NUM_BATCHES_PRECISE, len(precise_bn_loader)),
cfg.NUM_GPUS > 0,
)
_ = misc.aggregate_sub_bn_stats(model)
# Save a checkpoint.
if is_checkp_epoch:
cu.save_checkpoint(cfg.OUTPUT_DIR, model, optimizer, cur_epoch,
cfg)
# Evaluate the model on validation set.
if is_eval_epoch:
is_best_epoch, _ = eval_epoch(val_loader, model, val_meter,
cur_epoch, cfg, writer, wandb_log)
if is_best_epoch:
cu.save_checkpoint(cfg.OUTPUT_DIR,
model,
optimizer,
cur_epoch,
cfg,
is_best_epoch=is_best_epoch)
if writer is not None:
writer.close()
def train(cfg):
"""
Train a video model for many epochs on train set and evaluate it on val set.
Args:
cfg (CfgNode): configs. Details can be found in
slowfast/config/defaults.py
"""
# Set up environment.
du.init_distributed_training(cfg)
# Set random seed from configs.
np.random.seed(cfg.RNG_SEED)
torch.manual_seed(cfg.RNG_SEED)
# Setup logging format.
logging.setup_logging(cfg.OUTPUT_DIR)
# Init multigrid.
multigrid = None
if cfg.MULTIGRID.LONG_CYCLE or cfg.MULTIGRID.SHORT_CYCLE:
multigrid = MultigridSchedule()
cfg = multigrid.init_multigrid(cfg)
if cfg.MULTIGRID.LONG_CYCLE:
cfg, _ = multigrid.update_long_cycle(cfg, cur_epoch=0)
# Print config.
logger.info("Train with config:")
logger.info(pprint.pformat(cfg))
# Build the video model and print model statistics.
model = build_model(cfg)
if du.is_master_proc() and cfg.LOG_MODEL_INFO:
misc.log_model_info(model, cfg, use_train_input=True)
# Construct the optimizer.
optimizer = optim.construct_optimizer(model, cfg)
# Create a GradScaler for mixed precision training
scaler = torch.cuda.amp.GradScaler(enabled=cfg.TRAIN.MIXED_PRECISION)
# Load a checkpoint to resume training if applicable.
if cfg.TRAIN.AUTO_RESUME and cu.has_checkpoint(cfg.OUTPUT_DIR):
logger.info("Load from last checkpoint.")
last_checkpoint = cu.get_last_checkpoint(cfg.OUTPUT_DIR, task=cfg.TASK)
if last_checkpoint is not None:
checkpoint_epoch = cu.load_checkpoint(
last_checkpoint,
model,
cfg.NUM_GPUS > 1,
optimizer,
scaler if cfg.TRAIN.MIXED_PRECISION else None,
)
start_epoch = checkpoint_epoch + 1
elif "ssl_eval" in cfg.TASK:
last_checkpoint = cu.get_last_checkpoint(cfg.OUTPUT_DIR,
task="ssl")
checkpoint_epoch = cu.load_checkpoint(
last_checkpoint,
model,
cfg.NUM_GPUS > 1,
optimizer,
scaler if cfg.TRAIN.MIXED_PRECISION else None,
epoch_reset=True,
clear_name_pattern=cfg.TRAIN.CHECKPOINT_CLEAR_NAME_PATTERN,
)
start_epoch = checkpoint_epoch + 1
else:
start_epoch = 0
elif cfg.TRAIN.CHECKPOINT_FILE_PATH != "":
logger.info("Load from given checkpoint file.")
checkpoint_epoch = cu.load_checkpoint(
cfg.TRAIN.CHECKPOINT_FILE_PATH,
model,
cfg.NUM_GPUS > 1,
optimizer,
scaler if cfg.TRAIN.MIXED_PRECISION else None,
inflation=cfg.TRAIN.CHECKPOINT_INFLATE,
convert_from_caffe2=cfg.TRAIN.CHECKPOINT_TYPE == "caffe2",
epoch_reset=cfg.TRAIN.CHECKPOINT_EPOCH_RESET,
clear_name_pattern=cfg.TRAIN.CHECKPOINT_CLEAR_NAME_PATTERN,
)
start_epoch = checkpoint_epoch + 1
else:
start_epoch = 0
# Create the video train and val loaders.
train_loader = loader.construct_loader(cfg, "train")
val_loader = loader.construct_loader(cfg, "val")
precise_bn_loader = (loader.construct_loader(
cfg, "train", is_precise_bn=True)
if cfg.BN.USE_PRECISE_STATS else None)
# if (
# cfg.TASK == "ssl"
# and cfg.MODEL.MODEL_NAME == "ContrastiveModel"
# and cfg.CONTRASTIVE.KNN_ON
# ):
# if hasattr(model, "module"):
# model.module.init_knn_labels(train_loader)
# else:
# model.init_knn_labels(train_loader)
# Create meters.
if cfg.DETECTION.ENABLE:
train_meter = AVAMeter(len(train_loader), cfg, mode="train")
val_meter = AVAMeter(len(val_loader), cfg, mode="val")
else:
train_meter = TrainMeter(1e6, cfg)
val_meter = ValMeter(1e6, cfg)
# set up writer for logging to Tensorboard format.
if cfg.TENSORBOARD.ENABLE and du.is_master_proc(
cfg.NUM_GPUS * cfg.NUM_SHARDS):
writer = tb.TensorboardWriter(cfg)
else:
writer = None
# Perform the training loop.
logger.info("Start epoch: {}".format(start_epoch + 1))
epoch_timer = EpochTimer()
for cur_epoch in range(start_epoch, cfg.SOLVER.MAX_EPOCH):
if cur_epoch > 0 and cfg.DATA.LOADER_CHUNK_SIZE > 0:
num_chunks = math.ceil(cfg.DATA.LOADER_CHUNK_OVERALL_SIZE /
cfg.DATA.LOADER_CHUNK_SIZE)
skip_rows = (cur_epoch) % num_chunks * cfg.DATA.LOADER_CHUNK_SIZE
logger.info(
f"=================+++ num_chunks {num_chunks} skip_rows {skip_rows}"
)
cfg.DATA.SKIP_ROWS = skip_rows
logger.info(f"|===========| skip_rows {skip_rows}")
train_loader = loader.construct_loader(cfg, "train")
loader.shuffle_dataset(train_loader, cur_epoch)
if cfg.MULTIGRID.LONG_CYCLE:
cfg, changed = multigrid.update_long_cycle(cfg, cur_epoch)
if changed:
(
model,
optimizer,
train_loader,
val_loader,
precise_bn_loader,
train_meter,
val_meter,
) = build_trainer(cfg)
# Load checkpoint.
if cu.has_checkpoint(cfg.OUTPUT_DIR):
last_checkpoint = cu.get_last_checkpoint(cfg.OUTPUT_DIR,
task=cfg.TASK)
assert "{:05d}.pyth".format(cur_epoch) in last_checkpoint
else:
last_checkpoint = cfg.TRAIN.CHECKPOINT_FILE_PATH
logger.info("Load from {}".format(last_checkpoint))
cu.load_checkpoint(last_checkpoint, model, cfg.NUM_GPUS > 1,
optimizer)
# Shuffle the dataset.
loader.shuffle_dataset(train_loader, cur_epoch)
if hasattr(train_loader.dataset, "_set_epoch_num"):
train_loader.dataset._set_epoch_num(cur_epoch)
# Train for one epoch.
epoch_timer.epoch_tic()
train_epoch(
train_loader,
model,
optimizer,
scaler,
train_meter,
cur_epoch,
cfg,
writer,
)
epoch_timer.epoch_toc()
logger.info(
f"Epoch {cur_epoch} takes {epoch_timer.last_epoch_time():.2f}s. Epochs "
f"from {start_epoch} to {cur_epoch} take "
f"{epoch_timer.avg_epoch_time():.2f}s in average and "
f"{epoch_timer.median_epoch_time():.2f}s in median.")
logger.info(
f"For epoch {cur_epoch}, each iteraction takes "
f"{epoch_timer.last_epoch_time()/len(train_loader):.2f}s in average. "
f"From epoch {start_epoch} to {cur_epoch}, each iteraction takes "
f"{epoch_timer.avg_epoch_time()/len(train_loader):.2f}s in average."
)
is_checkp_epoch = (cu.is_checkpoint_epoch(
cfg,
cur_epoch,
None if multigrid is None else multigrid.schedule,
) or cur_epoch == cfg.SOLVER.MAX_EPOCH - 1)
is_eval_epoch = misc.is_eval_epoch(
cfg, cur_epoch, None if multigrid is None else multigrid.schedule)
# Compute precise BN stats.
if ((is_checkp_epoch or is_eval_epoch) and cfg.BN.USE_PRECISE_STATS
and len(get_bn_modules(model)) > 0):
calculate_and_update_precise_bn(
precise_bn_loader,
model,
min(cfg.BN.NUM_BATCHES_PRECISE, len(precise_bn_loader)),
cfg.NUM_GPUS > 0,
)
_ = misc.aggregate_sub_bn_stats(model)
# Save a checkpoint.
if is_checkp_epoch:
cu.save_checkpoint(
cfg.OUTPUT_DIR,
model,
optimizer,
cur_epoch,
cfg,
scaler if cfg.TRAIN.MIXED_PRECISION else None,
)
# Evaluate the model on validation set.
if is_eval_epoch:
eval_epoch(
val_loader,
model,
val_meter,
cur_epoch,
cfg,
train_loader,
writer,
)
if writer is not None:
writer.close()
result_string = "Top1 Acc: {:.2f} Top5 Acc: {:.2f} MEM: {:.2f}" "".format(
100 - val_meter.min_top1_err,
100 - val_meter.min_top5_err,
misc.gpu_mem_usage(),
)
logger.info("training done: {}".format(result_string))
return result_string