# Copyright (c) 2021 Graphcore Ltd. All rights reserved. import sys import argparse from datetime import datetime import logging import math import numpy as np from pathlib import Path import os import time from tqdm import tqdm from typing import Callable, Union import yacs import wandb import poptorch from poptorch import trainingModel, inferenceModel, DataLoader, DataLoaderMode from poptorch.optim import Optimizer, SGD import torch import torch.optim.lr_scheduler as lr_scheduler from torch.utils.data import DataLoader as torchDataLoader from torchinfo import summary from models.detector import Detector from models.yolov4_p5 import Yolov4P5 from models.recomputation import recomputation_checkpoint from utils.config import get_cfg_defaults, override_cfg, save_cfg from utils.dataset import Dataset from utils.parse_args import parse_args from utils.postprocessing import post_processing, IPUPredictionsPostProcessing from utils.tools import load_and_fuse_pretrained_weights, StatRecorder from utils.visualization import plotting_tool from utils.weight_avg import average_model_weights path_to_detection = Path(__file__).parent.resolve() os.environ["PYTORCH_APPS_DETECTION_PATH"] = str(path_to_detection) logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO) logger = logging.getLogger("Detector") def ipu_options(cfg: yacs.config.CfgNode, model: Detector, mode: str, opt: argparse.ArgumentParser = None): """Configurate the IPU options using cfg and opt options. Parameters: opt: opt object containing options introduced in the command line cfg: yacs object containing the config model[Detector]: a torch Detector Model mode: str indicating 'train', 'test' or 'test_inference' Returns: ipu_opts: Options for the IPU configuration """ device_iterations = cfg.ipuopts.device_iterations ipu_opts = poptorch.Options() ipu_opts.deviceIterations(device_iterations) ipu_opts.autoRoundNumIPUs(True) # Calculate the number of replicas from the pipeline length ipu_opts.replicationFactor(cfg.system.num_ipus // (1 + len(cfg.model.pipeline_splits))) ipu_opts.enableExecutableCaching(cfg.training.exec_cache_path) # Compile offline (no IPUs required) if opt is not None and opt.compile_only: ipu_opts.useOfflineIpuTarget() if opt is not None and opt.profile_dir: ipu_opts.enableProfiling(opt.profile_dir) if cfg.ipuopts.available_memory_proportion: amp = cfg.ipuopts.available_memory_proportion if isinstance(amp, float): amp_dict = {f"IPU{i}": amp for i in range(cfg.system.num_ipus)} elif isinstance(cfg.ipuopts.available_memory_proportion, list): assert len(amp) == len(cfg.model.pipeline_splits) + 1 amp_dict = {f"IPU{i}": value for i, value in enumerate(amp)} else: raise TypeError("Wrong type of cfg.ipuopts.available_memory_proportion. " "Use either float or list.") ipu_opts.setAvailableMemoryProportion(amp_dict) if opt is not None and opt.benchmark: ipu_opts.Distributed.disable() if cfg.model.precision == "half": ipu_opts.Precision.setPartialsType(torch.float16) model.half() elif cfg.model.precision == "mixed": ipu_opts.Precision.setPartialsType(torch.float16) model.half() model.headp3 = model.headp3.float() model.headp4 = model.headp4.float() model.headp5 = model.headp5.float() elif cfg.model.precision != "single": raise ValueError("Only supoprt half, mixed or single precision") if mode == "train": ipu_opts.Training.gradientAccumulation(cfg.ipuopts.gradient_accumulation) ipu_opts.outputMode(poptorch.OutputMode.Sum) ipu_opts.Training.setAutomaticLossScaling(enabled=cfg.training.auto_loss_scaling) ipu_opts.Precision.enableStochasticRounding(cfg.training.stochastic_rounding) if cfg.model.sharded: ipu_opts.setExecutionStrategy(poptorch.ShardedExecution()) else: ipu_opts.setExecutionStrategy(poptorch.PipelinedExecution(poptorch.AutoStage.AutoIncrement)) return ipu_opts def get_loader(opt: argparse.ArgumentParser, cfg: yacs.config.CfgNode, ipu_opts: poptorch.Options, mode: str): """Gets a new loader for the model. Parameters: opt: opt object containing options introduced in the command line cfg: yacs object containing the config ipu_opts: Options for the IPU configuration mode: str indicating 'train', 'test' or 'test_inference' Returns: model[Detector]: a torch Detector Model """ dataset = Dataset(path=opt.data, cfg=cfg, mode=mode) shuffle = mode == "train" # Creates a loader using the dataset if cfg.model.ipu: loader = DataLoader( ipu_opts, dataset, batch_size=cfg.model.micro_batch_size, shuffle=shuffle, num_workers=cfg.system.num_workers, mode=DataLoaderMode.Async, ) else: loader = torchDataLoader(dataset, batch_size=cfg.model.micro_batch_size, shuffle=shuffle) return loader def get_optimizer(cfg: yacs.config.CfgNode, model: Detector): """Returns the correct type of optimizer for ipu/cpu training Parameters: cfg: yacs object containing the config model: a Detector object Returns: optimizer: a torch/poptorch optimizer """ nominal_batch_size = 64 weight_decay = ( cfg.training.weight_decay * cfg.model.micro_batch_size * cfg.ipuopts.gradient_accumulation / nominal_batch_size ) param_group0, param_group1, param_group2 = [], [], [] for k, v in model.named_parameters(): v.requires_grad = True if ".bias" in k: param_group2.append(v) elif ".weight" in k and ".norm" not in k: param_group1.append(v) else: param_group0.append(v) if cfg.model.ipu: optimizer = SGD( param_group0, lr=cfg.training.initial_lr, momentum=cfg.training.momentum, loss_scaling=cfg.training.loss_scaling, use_combined_accum=True, dampening=0, ) else: optimizer = SGD(param_group0, lr=cfg.training.initial_lr, momentum=cfg.training.momentum) optimizer.add_param_group({"params": param_group1, "weight_decay": weight_decay}) optimizer.add_param_group({"params": param_group2}) return optimizer def get_model_and_loader(opt: argparse.ArgumentParser, cfg: yacs.config.CfgNode, mode: str): """Prepares the model and gets a new loader for the model. Parameters: opt: opt object containing options introduced in the command line cfg: yacs object containing the config mode: str indicating 'train', 'test' or 'test_inference' Returns: model[Detector]: a torch Detector Model loader[DataLoader]: a torch or poptorch DataLoader containing the specified dataset on "cfg" """ # Create model model = Yolov4P5(cfg) # Insert the pipeline splits if using pipeline if cfg.model.pipeline_splits: named_layers = {name: layer for name, layer in model.named_modules()} for ipu_idx, split in enumerate(cfg.model.pipeline_splits): named_layers[split] = poptorch.BeginBlock(ipu_id=ipu_idx + 1, layer_to_call=named_layers[split]) if len(cfg.model.recomputation_ckpts): for name, layer in model.named_modules(): if name in cfg.model.recomputation_ckpts: recomputation_checkpoint(layer) # Load weights and fuses some batch normalizations with some convolutions if cfg.model.normalization == "batch": if opt.weights: print("loading pretrained weights") model = load_and_fuse_pretrained_weights(model, opt.weights, mode != "train") if mode == "train": model.train() else: model.optimize_for_inference() model.eval() if opt.print_summary: summary( model, input_size=( cfg.model.micro_batch_size, cfg.model.input_channels, cfg.model.image_size, cfg.model.image_size, ), ) print("listing all layers by names") named_layers = {name: layer for name, layer in model.named_modules()} for layer in named_layers: print(layer) # Create the specific ipu options if cfg.model.ipu ipu_opts = ipu_options(cfg, model, mode, opt) if cfg.model.ipu else None # Creates the loader loader = get_loader(opt, cfg, ipu_opts, mode) # Calls the poptorch wrapper and compiles the model if cfg.model.ipu: img, labels, _, _ = next(iter(loader)) if cfg.model.mode == "train": optimizer = get_optimizer(cfg, model) model = trainingModel(model, ipu_opts, optimizer=optimizer) else: model = inferenceModel(model, ipu_opts) model.compile(img) if opt.benchmark: warm_up_iterations = 100 for _ in range(warm_up_iterations): _ = model(img) return model, loader def set_warmup_lr_and_momentum( optimizer: Optimizer, num_warmup: int, batch_idx: int, lr_lambda: Callable[[int], float], epoch: int, momentum: float, ): xi = [0, num_warmup] for j, x in enumerate(optimizer.param_groups): x["lr"] = np.interp(batch_idx, xi, [0.1 if j == 2 else 0.0, x["initial_lr"] * lr_lambda(epoch + 1)]) if "momentum" in x: x["momentum"] = np.interp(batch_idx, xi, [0.9, momentum]) return optimizer def inference( opt: argparse.ArgumentParser, cfg: yacs.config.CfgNode, model: Union[Detector, poptorch.PoplarExecutor], loader: Union[torchDataLoader, DataLoader], stat_recorder: StatRecorder, run_coco_eval: bool = False, ): inference_progress = tqdm(loader) inference_progress.set_description("Running inference") stat_recorder.reset_eval_stats() for batch_idx, (transformed_images, transformed_labels, image_sizes, image_indxs) in enumerate(inference_progress): start_time = time.time() y = model(transformed_images) inference_step_time = time.time() - start_time inference_round_trip_time = ( model.getLatency() if cfg.model.ipu else (inference_step_time,) * 3 ) # returns (min, max, avg) latency processed_batch = post_processing(cfg, y, image_sizes, transformed_labels) stat_recorder.record_inference_stats(inference_round_trip_time, inference_step_time) if cfg.inference.plot_output and batch_idx % cfg.inference.plot_step == 0: img_paths = plotting_tool( cfg, processed_batch[0], [loader.dataset.get_image(img_idx) for img_idx in image_indxs] ) if opt.wandb: wandb.log({"inference_batch_{}".format(batch_idx): [wandb.Image(path) for path in img_paths]}) if opt.benchmark and batch_idx == 100: break pruned_preds_batch = processed_batch[0] processed_labels_batch = processed_batch[1] if cfg.eval.metrics: for idx, (pruned_preds, processed_labels) in enumerate(zip(pruned_preds_batch, processed_labels_batch)): stat_recorder.record_eval_stats( processed_labels, pruned_preds, image_sizes[idx], loader.dataset.images_id[image_indxs[idx]], run_coco_eval, ) stat_recorder.logging(print, run_coco_eval) if __name__ == "__main__": opt = parse_args() if len(opt.data) > 0.0 and opt.data[-1] != "/": opt.data += "/" cfg = get_cfg_defaults() cfg.merge_from_file(opt.config) cfg = override_cfg(opt, cfg) cfg.freeze() config_filename = Path(opt.config) config_filename = config_filename.with_name(f"override-{config_filename.name}") save_cfg(config_filename, cfg) if opt.show_config: logger.info(f"Model options: \n'{cfg}'") stat_recorder = StatRecorder(cfg, opt.data, opt.wandb) model, loader = get_model_and_loader(opt, cfg, cfg.model.mode) if opt.compile_only: logger.info("Model successfully compiled. Exiting now as '--compile-only' argument was passed.") sys.exit(0) if cfg.model.mode == "train": logger.error("Training to be implemented!") else: run_coco_eval = cfg.eval.metrics and not opt.benchmark inference(opt, cfg, model, loader, stat_recorder, run_coco_eval)