def evaluate(model_path: Path, datasets: typing.List[Datasets], output_folder: Path, find_mistakes: bool = False, include_heading: bool = False) -> str: """Evaluate a model, returning the results as CSV. Args: model_path (Path): path to the model folder containing the YAML file and the saved weights datasets (typing.List[Datasets]): the datasets to evaluate on output_folder (Path): output folder for the mistake images (if applicable) find_mistakes (bool, optional): whether to output all mistakes as images to the output folder. Defaults to False. include_heading (bool, optional): whether to include a heading in the CSV output. Defaults to False. Raises: ValueError: if the YAML config file is missing Returns: str: the CSV string """ model_name = model_path.stem config_file = model_path.parent / f"{model_name}.yaml" if not config_file.exists(): raise ValueError("config file missing") cfg = CN.load_yaml_with_base(config_file) model = torch.load(model_path, map_location=DEVICE) model = device(model) model.eval() datasets = {mode: build_dataset(cfg, mode) for mode in datasets} classes = next(iter(datasets.values())).classes csv = [] if include_heading: csv.append(_csv_heading(classes)) for mode, dataset in datasets.items(): # Load dataset loader = build_data_loader(cfg, dataset, mode) # Compute statistics over whole dataset agg = StatsAggregator(classes) for images, labels in device(loader): predictions = model(images) agg.add_batch(predictions, labels, **(dict(inputs=images) if find_mistakes else dict())) csv.append(_csv(model, agg, model_name, mode)) if find_mistakes: groundtruth, mistakes = zip(*sorted(agg.mistakes, key=lambda x: x[0])) imgs = torch.tensor(mistakes).permute((0, 2, 3, 1)) imgs = unnormalize(imgs).permute((0, 3, 1, 2)) img = torchvision.utils.make_grid(imgs, pad_value=1, nrow=4) img = img.numpy().transpose((1, 2, 0)) * 255 img = Image.fromarray(img.astype(np.uint8)) mistakes_file = output_folder / \ f"{model_name}_{mode.value}_mistakes.png" logger.info(f"Writing mistakes to {mistakes_file}") img.save(mistakes_file) groundtruth_file = output_folder / \ f"{model_name}_{mode.value}_groundtruth.csv" with groundtruth_file.open("w") as f: f.write(",".join(map(str, groundtruth))) return "\n".join(csv)
def aggregator() -> StatsAggregator: agg = StatsAggregator(["a", "b"]) a_output = np.array([.9, .1, .8, .2, .9, .9, .9, .2]) b_output = 1 - a_output outputs = torch.tensor(np.stack([a_output, b_output], axis=-1)) labels = torch.tensor([0, 0, 0, 1, 0, 0, 1, 0]) # predicted: [0, 1, 0, 1, 0, 0, 0, 1] agg.add_batch(outputs, labels) return agg
def test_f1_score(aggregator: StatsAggregator): a_precision = 4 / 5 a_recall = 4 / 6 assert np.isclose(aggregator.f1_score("a"), 2 * a_precision * a_recall, a_precision + a_recall) b_precision = 1 / 3 b_recall = 1 / 2 assert np.isclose(aggregator.f1_score("b"), 2 * b_precision * b_recall, b_precision + b_recall)
def _csv(model: torch.nn.Module, agg: StatsAggregator, model_name: str, mode: Datasets) -> str: params = sum([np.prod(p.size()) for p in model.parameters()]) return ",".join(map(str, [model_name, mode.value, params, agg.accuracy(), *map(agg.precision, agg.classes), *map(agg.recall, agg.classes), *map(agg.f1_score, agg.classes), *agg.confusion_matrix.flatten() ]))
def test_empty_batch(): aggregator = StatsAggregator(["a", "b"]) assert aggregator.accuracy() == 0 assert aggregator.precision("a") == 0 assert aggregator.precision("b") == 0 assert aggregator.recall("a") == 0 assert aggregator.recall("b") == 0 assert aggregator.f1_score("a") == 0 assert aggregator.f1_score("b") == 0
def test_recall(aggregator: StatsAggregator): assert np.isclose(aggregator.recall("a"), 4 / 6) assert np.isclose(aggregator.recall("b"), 1 / 2)
def test_precision(aggregator: StatsAggregator): assert np.isclose(aggregator.precision("a"), 4 / 5) assert np.isclose(aggregator.precision("b"), 1 / 3)
def test_accuracy(aggregator: StatsAggregator): assert np.isclose(aggregator.accuracy(), 5 / 8)
def train_model(cfg: CN, run_dir: Path, model: torch.nn.Module, is_inception: bool = False, model_name: str = None, eval_on_train: bool = False) -> nn.Module: """Train a model that has already been loaded. Args: cfg (CN): the configuration object describing the model, dataset, etc. run_dir (Path): where to write tensorboard files, the active YAML file, and the chosen weights model (torch.nn.Module): the loaded model is_inception (bool, optional): whether the model is InceptionV3. Defaults to False. model_name (str, optional): the name of the model (by default the last component of the run directory). Defaults to None. eval_on_train (bool, optional): whether to evaluate on the training set. Defaults to False. Returns: nn.Module: the trained model """ logger.info(f"Starting training in {run_dir}") if not model_name: model_name = run_dir.name # Create folder if run_dir.exists(): logger.warning( f"The folder {run_dir} already exists and will be overwritten by this run" ) shutil.rmtree(run_dir, ignore_errors=True) run_dir.mkdir(parents=True, exist_ok=True) # Store config with (run_dir / f"{model_name}.yaml").open("w") as f: cfg.dump(stream=f) # Move model to device device(model) best_weights, best_accuracy, best_step = copy.deepcopy( model.state_dict()), 0., 0 criterion = nn.CrossEntropyLoss() modes = {Datasets.TRAIN, Datasets.VAL} if eval_on_train: dataset = build_dataset(cfg, Datasets.TRAIN) datasets = {mode: dataset for mode in modes} else: datasets = {mode: build_dataset(cfg, mode) for mode in modes} classes = datasets[Datasets.TRAIN].classes loader = { mode: build_data_loader(cfg, datasets[mode], mode) for mode in modes } writer = {mode: SummaryWriter(run_dir / mode.value) for mode in modes} aggregator = {mode: StatsAggregator(classes) for mode in modes} def log(step: int, loss: float, mode: Datasets): if mode == Datasets.TRAIN: logger.info(f"Step {step:5d}: loss {loss:.3f}") w, agg = (x[mode] for x in (writer, aggregator)) w.add_scalar("Loss", loss, step) w.add_scalar("Accuracy", agg.accuracy(), step) for c in classes: w.add_scalar(f"Precision/{c}", agg.precision(c), step) w.add_scalar(f"Recall/{c}", agg.recall(c), step) w.add_scalar(f"F1 score/{c}", agg.f1_score(c), step) def perform_iteration(data: typing.Tuple[torch.Tensor, torch.Tensor], mode: Datasets): inputs, labels = map(device, data) with torch.set_grad_enabled(mode == Datasets.TRAIN): # Reset gradients optimizer.zero_grad() # Forward pass and compute loss if is_inception and mode == Datasets.TRAIN: # Special case for inception models outputs, auxiliary_outputs = model(inputs) loss1 = criterion(outputs, labels) loss2 = criterion(auxiliary_outputs, labels) loss = loss1 + 0.4 * loss2 else: outputs = model(inputs) loss = criterion(outputs, labels) if mode == Datasets.TRAIN: loss.backward() with torch.no_grad(): aggregator[mode].add_batch(outputs, labels) # Perform optimisation if mode == Datasets.TRAIN: optimizer.step() # Return return loss.item() step = 0 log_every_n = 100 # Ensure we're in training mode model.train() # Loop over training phases for phase in cfg.TRAINING.PHASES: for p in model.parameters(): p.requires_grad = False parameters = list(model.parameters()) if phase.PARAMS == "all" \ else model.params[phase.PARAMS] for p in parameters: p.requires_grad = True optimizer = build_optimizer_from_config(phase.OPTIMIZER, parameters) # Loop over epochs (passes over the whole dataset) for epoch in range(phase.EPOCHS): aggregator[Datasets.TRAIN].reset() # Iterate the training set losses = [] for i, data in enumerate(loader[Datasets.TRAIN]): # Perform training iteration losses.append(perform_iteration(data, mode=Datasets.TRAIN)) if step % log_every_n == 0: loss = np.mean(list(losses)) log(step, loss, Datasets.TRAIN) aggregator[Datasets.TRAIN].reset() losses = [] # Validate entire validation dataset model.eval() aggregator[Datasets.VAL].reset() # Iterate entire val dataset perform_val_iteration = functools.partial( perform_iteration, mode=Datasets.VAL) val_losses = map(perform_val_iteration, loader[Datasets.VAL]) # Gather losses and log val_loss = np.mean(list(val_losses)) log(step, val_loss, Datasets.VAL) model.train() # Save weights if we get a better performance accuracy = aggregator[Datasets.VAL].accuracy() if accuracy >= best_accuracy: best_accuracy = accuracy best_weights = copy.deepcopy(model.state_dict()) best_step = step # Get ready for next step step += 1 # Clean up for w in writer.values(): w.flush() w.close() logger.info("Finished training") logger.info( f"Restoring best weight state (step {best_step} with validation accuracy of {best_accuracy})" ) model.load_state_dict(best_weights) torch.save(model, run_dir / f"{model_name}.pt") with (run_dir / f"{model_name}.txt").open("w") as f: f.write(f"exported at step: {best_step}") return model