def generate_hierarchy_vis(args): path = get_graph_path_from_args(**vars(args)) print('==> Reading from {}'.format(path)) G = read_graph(path) roots = list(get_roots(G)) num_roots = len(roots) root = args.vis_root or next(get_roots(G)) assert root in G, f'Node {root} is not a valid node. Nodes: {G.nodes}' dataset = None if args.dataset and args.vis_leaf_images: cls = getattr(data, args.dataset) dataset = cls(root='./data', train=False, download=True) color_info = get_color_info(G, args.color, color_leaves=not args.vis_no_color_leaves, color_path_to=args.vis_color_path_to, color_nodes=args.vis_color_nodes or ()) node_to_conf = generate_node_conf(args.vis_node_conf) tree = build_tree(G, root, color_info=color_info, force_labels_left=args.vis_force_labels_left or [], dataset=dataset, include_leaf_images=args.vis_leaf_images, image_resize_factor=args.vis_image_resize_factor, include_fake_sublabels=args.vis_fake_sublabels, node_to_conf=node_to_conf) graph = build_graph(G) if num_roots > 1: Colors.red(f'Found {num_roots} roots! Should be only 1: {roots}') else: print(f'Found just {num_roots} root.') fname = generate_vis_fname(**vars(args)) parent = Path(fwd()).parent generate_vis(str(parent / 'nbdt/templates/tree-template.html'), tree, fname, zoom=args.vis_zoom, straight_lines=not args.vis_curved, show_sublabels=args.vis_sublabels, height=args.vis_height, width=args.vis_width, dark=args.vis_dark, margin_top=args.vis_margin_top, margin_left=args.vis_margin_left, hide=args.vis_hide or [], above_dy=args.vis_above_dy, below_dy=args.vis_below_dy, scale=args.vis_scale, root_y=args.vis_root_y, colormap=args.vis_colormap)
def print_stats(leaves_seen, wnid_set, tree_name, node_type): print( f"[{tree_name}] \t {node_type}: {len(leaves_seen)} \t WNIDs missing from {node_type}: {len(wnid_set)}" ) if len(wnid_set): Colors.red( f"==> Warning: WNIDs in wnid.txt are missing from {tree_name} {node_type}" )
def f(**kwargs): try: net = init(**optional_kwargs, **kwargs) except TypeError as e: # likely because `dataset` not allowed arg print(e) try: net = init(**kwargs) except Exception as e: Colors.red(f"Fatal error: {e}") exit() return net
def test_hierarchy(args): wnids = get_wnids_from_dataset(args.dataset) path = get_graph_path_from_args(**vars(args)) print("==> Reading from {}".format(path)) G = read_graph(path) G_name = Path(path).stem leaves_seen, wnid_set1 = match_wnid_leaves(wnids, G, G_name) print_stats(leaves_seen, wnid_set1, G_name, "leaves") leaves_seen, wnid_set2 = match_wnid_nodes(wnids, G, G_name) print_stats(leaves_seen, wnid_set2, G_name, "nodes") num_roots = len(list(get_roots(G))) if num_roots == 1: Colors.green("Found just 1 root.") else: Colors.red(f"Found {num_roots} roots. Should be only 1.") if len(wnid_set1) == len(wnid_set2) == 0 and num_roots == 1: Colors.green("==> All checks pass!") else: Colors.red("==> Test failed")
def generate_hierarchy_vis(args): path = get_graph_path_from_args(**vars(args)) print('==> Reading from {}'.format(path)) G = read_graph(path) roots = list(get_roots(G)) num_roots = len(roots) root = next(get_roots(G)) dataset = None if args.dataset: cls = getattr(data, args.dataset) dataset = cls(root='./data', train=False, download=True) color_info = get_color_info(G, args.color, color_leaves=not args.vis_no_color_leaves, color_path_to=args.vis_color_path_to, color_nodes=args.vis_color_nodes or ()) tree = build_tree(G, root, color_info=color_info, force_labels_left=args.vis_force_labels_left or [], dataset=dataset, include_leaf_images=args.vis_leaf_images, image_resize_factor=args.vis_image_resize_factor) graph = build_graph(G) if num_roots > 1: Colors.red(f'Found {num_roots} roots! Should be only 1: {roots}') else: print(f'Found just {num_roots} root.') fname = generate_vis_fname(**vars(args)) parent = Path(fwd()).parent generate_vis(str(parent / 'nbdt/templates/tree-template.html'), tree, 'tree', fname, zoom=args.vis_zoom, straight_lines=not args.vis_curved, show_sublabels=args.vis_sublabels, height=args.vis_height, dark=args.vis_dark)
def generate_vis(path_template, data, name, fname, zoom=2, straight_lines=True, show_sublabels=False, height=750, dark=False): with open(path_template) as f: html = f.read() \ .replace( "CONFIG_TREE_DATA", json.dumps([data])) \ .replace( "CONFIG_ZOOM", str(zoom)) \ .replace( "CONFIG_STRAIGHT_LINES", str(straight_lines).lower()) \ .replace( "CONFIG_SHOW_SUBLABELS", str(show_sublabels).lower()) \ .replace( "CONFIG_TITLE", fname) \ .replace( "CONFIG_VIS_HEIGHT", str(height)) \ .replace( "CONFIG_BG_COLOR", "#111111" if dark else "#FFFFFF") \ .replace( "CONFIG_TEXT_COLOR", '#FFFFFF' if dark else '#000000') \ .replace( "CONFIG_TEXT_RECT_COLOR", "rgba(17,17,17,0.8)" if dark else "rgba(255,255,255,0.8)") os.makedirs('out', exist_ok=True) path_html = f'out/{fname}-{name}.html' with open(path_html, 'w') as f: f.write(html) Colors.green('==> Wrote HTML to {}'.format(path_html))
def test(epoch, checkpoint=True): nonlocal best_acc net.eval() test_loss = 0 metric.clear() with torch.no_grad(): for batch_idx, (inputs, targets) in enumerate(testloader): inputs, targets = inputs.to(device), targets.to(device) outputs = net(inputs) if not args.disable_test_eval: loss = criterion(outputs, targets) test_loss += loss.item() metric.forward(outputs, targets) transform = testset.transform_val_inverse().to(device) stat = analyzer.update_batch(outputs, targets, transform(inputs)) progress_bar( batch_idx, len(testloader), "Loss: %.3f | Acc: %.3f%% (%d/%d) %s" % ( test_loss / (batch_idx + 1), 100.0 * metric.report(), metric.correct, metric.total, f"| {analyzer.name}: {stat}" if stat else "", ), ) # Save checkpoint. acc = 100.0 * metric.report() print("Accuracy: {}, {}/{} | Best Accurracy: {}".format( acc, metric.correct, metric.total, best_acc)) if acc > best_acc and checkpoint: Colors.green(f"Saving to {checkpoint_fname} ({acc})..") state = { "net": net.state_dict(), "acc": acc, "epoch": epoch, } os.makedirs("checkpoint", exist_ok=True) torch.save(state, f"./checkpoint/{checkpoint_fname}.pth") best_acc = acc
def load_state_dict_from_key( keys, model_urls, pretrained=False, progress=True, root=".cache/torch/checkpoints", device="cpu", ): valid_keys = [key for key in keys if key in model_urls] if not valid_keys: raise UserWarning(f"None of the keys {keys} correspond to a pretrained model.") key = valid_keys[-1] url = model_urls[key] Colors.green(f"Loading pretrained model {key} from {url}") return load_state_dict_from_url( url, Path.home() / root, progress=progress, check_hash=False, map_location=torch.device(device), )
def __init__( self, *args, superclass_wnids, dataset_test=None, Rules=SoftRules, metric=None, **kwargs, ): """Pass wnids to classify. Assumes index of each wnid is the index of the wnid in the rules.wnids list. This agrees with Node.wnid_to_class_index as of writing, since rules.wnids = get_wnids(...). """ # TODO: for now, ignores metric super().__init__(*args, **kwargs) kwargs["dataset"] = dataset_test kwargs.pop("path_graph", "") kwargs.pop("path_wnids", "") self.rules_test = Rules(*args, **kwargs) self.superclass_wnids = superclass_wnids self.total = self.correct = 0 self.mapping_target, self.new_to_old_classes_target = Superclass.build_mapping( self.rules_test.tree.wnids_leaves, superclass_wnids) self.mapping_pred, self.new_to_old_classes_pred = Superclass.build_mapping( self.rules.tree.wnids_leaves, superclass_wnids) mapped_classes = [ self.classes[i] for i in (self.mapping_target >= 0).nonzero() ] Colors.cyan( f"==> Mapped {len(mapped_classes)} classes to your superclasses: " f"{mapped_classes}")
def generate_hierarchy( dataset, method, seed=0, branching_factor=2, extra=0, no_prune=False, fname="", path="", single_path=False, induced_linkage="ward", induced_affinity="euclidean", checkpoint=None, arch=None, model=None, **kwargs, ): wnids = get_wnids_from_dataset(dataset) if method == "wordnet": G = build_minimal_wordnet_graph(wnids, single_path) elif method == "random": G = build_random_graph(wnids, seed=seed, branching_factor=branching_factor) elif method == "induced": G = build_induced_graph( wnids, dataset=dataset, checkpoint=checkpoint, model=arch, linkage=induced_linkage, affinity=induced_affinity, branching_factor=branching_factor, state_dict=model.state_dict() if model is not None else None, ) else: raise NotImplementedError(f'Method "{method}" not yet handled.') print_graph_stats(G, "matched") assert_all_wnids_in_graph(G, wnids) if not no_prune: G = prune_single_successor_nodes(G) print_graph_stats(G, "pruned") assert_all_wnids_in_graph(G, wnids) if extra > 0: G, n_extra, n_imaginary = augment_graph(G, extra, True) print(f"[extra] \t Extras: {n_extra} \t Imaginary: {n_imaginary}") print_graph_stats(G, "extra") assert_all_wnids_in_graph(G, wnids) path = get_graph_path_from_args( dataset=dataset, method=method, seed=seed, branching_factor=branching_factor, extra=extra, no_prune=no_prune, fname=fname, path=path, single_path=single_path, induced_linkage=induced_linkage, induced_affinity=induced_affinity, checkpoint=checkpoint, arch=arch, ) write_graph(G, path) Colors.green("==> Wrote tree to {}".format(path)) return path
def generate_hierarchy_vis_from(G, dataset, path_html, color="blue", vis_root=None, vis_no_color_leaves=False, vis_color_path_to=None, vis_color_nodes=(), vis_theme="regular", vis_force_labels_left=(), vis_leaf_images=False, vis_image_resize_factor=1, vis_fake_sublabels=False, vis_zoom=2, vis_curved=False, vis_sublabels=False, vis_height=750, vis_width=1000, vis_margin_top=20, vis_margin_left=250, vis_hide=(), vis_above_dy=325, vis_below_dy=475, vis_scale=1, vis_root_y="null", vis_colormap="colormap_annotated.png", vis_node_conf=(), verbose=False, **kwargs): """ :param path_html: Where to write final hierarchy """ roots = list(get_roots(G)) num_roots = len(roots) root = vis_root or next(get_roots(G)) assert root in G, f"Node {root} is not a valid node. Nodes: {G.nodes}" color_info = get_color_info( G, color, color_leaves=not vis_no_color_leaves, color_path_to=vis_color_path_to, color_nodes=vis_color_nodes or (), theme=vis_theme, ) node_to_conf = generate_node_conf(vis_node_conf) tree = build_tree( G, root, color_info=color_info, force_labels_left=vis_force_labels_left or [], dataset=dataset, include_leaf_images=vis_leaf_images, image_resize_factor=vis_image_resize_factor, include_fake_sublabels=vis_fake_sublabels, node_to_conf=node_to_conf, ) graph = build_graph(G) if num_roots > 1: Colors.red(f"Found {num_roots} roots! Should be only 1: {roots}") elif verbose: print(f"Found just {num_roots} root.") parent = Path(fwd()).parent generate_vis( str(parent / "nbdt/templates/tree-template.html"), tree, path_html, zoom=vis_zoom, straight_lines=not vis_curved, show_sublabels=vis_sublabels, height=vis_height, bg=color_info["bg"], text_rect=color_info["text_rect"], width=vis_width, margin_top=vis_margin_top, margin_left=vis_margin_left, hide=vis_hide or [], above_dy=vis_above_dy, below_dy=vis_below_dy, scale=vis_scale, root_y=vis_root_y, colormap=vis_colormap, verbose=verbose, )
def generate_vis( path_template, data, path_html, zoom=2, straight_lines=True, show_sublabels=False, height=750, margin_top=20, above_dy=325, y_node_sep=170, hide=[], _print=False, out_dir=".", scale=1, colormap="colormap_annotated.png", below_dy=475, root_y="null", width=1000, margin_left=250, bg="#FFFFFF", text_rect="rgba(255,255,255,0.8)", stroke_width=0.45, verbose=False, ): fname = Path(path_html).stem out_dir = Path(path_html).parent with open(path_template) as f: html = ( f.read().replace("CONFIG_MARGIN_LEFT", str(margin_left)).replace( "CONFIG_VIS_WIDTH", str(width)).replace("CONFIG_SCALE", str(scale)).replace( "CONFIG_PRINT", str(_print).lower()).replace( "CONFIG_HIDE", str(hide)).replace( "CONFIG_Y_NODE_SEP", str(y_node_sep)).replace( "CONFIG_ABOVE_DY", str(above_dy)).replace( "CONFIG_BELOW_DY", str(below_dy)).replace( "CONFIG_TREE_DATA", json.dumps([data])).replace( "CONFIG_ZOOM", str(zoom)).replace( "CONFIG_STRAIGHT_LINES", str(straight_lines).lower()). replace( "CONFIG_SHOW_SUBLABELS", str(show_sublabels).lower()).replace( "CONFIG_TITLE", fname).replace("CONFIG_VIS_HEIGHT", str(height)).replace( "CONFIG_BG_COLOR", bg).replace( "CONFIG_TEXT_RECT_COLOR", text_rect).replace( "CONFIG_STROKE_WIDTH", str(stroke_width)).replace( "CONFIG_MARGIN_TOP", str(margin_top)).replace( "CONFIG_ROOT_Y", str(root_y)).replace( "CONFIG_COLORMAP", f"""<img src="{colormap}" style=" position: absolute; top: 40px; left: 80px; height: 250px; border: 4px solid #ccc;">""" if isinstance(colormap, str) and os.path.exists(colormap) else "", )) os.makedirs(out_dir, exist_ok=True) with open(path_html, "w") as f: f.write(html) if verbose: Colors.green("==> Wrote HTML to {}".format(path_html))
download=True, transform=transform_test) assert trainset.classes == testset.classes, (trainset.classes, testset.classes) trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=2) testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2) Colors.cyan( f'Training with dataset {args.dataset} and {len(trainset.classes)} classes' ) # Model print('==> Building model..') model = getattr(models, args.arch) model_kwargs = {'num_classes': len(trainset.classes)} if args.pretrained: print('==> Loading pretrained model..') try: net = model(pretrained=True, dataset=args.dataset, **model_kwargs) except TypeError as e: # likely because `dataset` not allowed arg print(e) try:
dataset = getattr(data, args.dataset) dataset_kwargs = generate_kwargs(args, dataset, name=f'Dataset {args.dataset}', keys=data.custom.keys, globals=globals()) trainset = dataset(**dataset_kwargs, root='./data', train=True, download=True, transform=transform_train) testset = dataset(**dataset_kwargs, root='./data', train=False, download=True, transform=transform_test) assert trainset.classes == testset.classes, (trainset.classes, testset.classes) trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=2) testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2) Colors.cyan(f'Training with dataset {args.dataset} and {len(trainset.classes)} classes') # Model print('==> Building model..') model = getattr(models, args.arch) model_kwargs = {'num_classes': len(trainset.classes) } if args.pretrained: print('==> Loading pretrained model..') net = model(pretrained=True, dataset=args.dataset, **model_kwargs) else: print('==> Loading NBDT model..') net = ResNet18() net = SoftNBDT(pretrained=True, dataset='CIFAR10', arch='ResNet18', model=net) #Analyzer
def generate_vis(path_template, data, fname, zoom=2, straight_lines=True, show_sublabels=False, height=750, dark=False, margin_top=20, above_dy=325, y_node_sep=170, hide=[], _print=False, out_dir='.', scale=1, colormap='colormap_annotated.png', below_dy=475, root_y='null', width=1000, margin_left=250): with open(path_template) as f: html = f.read() \ .replace( "CONFIG_MARGIN_LEFT", str(margin_left)) \ .replace( "CONFIG_VIS_WIDTH", str(width)) \ .replace( "CONFIG_SCALE", str(scale)) \ .replace( "CONFIG_PRINT", str(_print).lower()) \ .replace( "CONFIG_HIDE", str(hide)) \ .replace( "CONFIG_Y_NODE_SEP", str(y_node_sep)) \ .replace( "CONFIG_ABOVE_DY", str(above_dy)) \ .replace( "CONFIG_BELOW_DY", str(below_dy)) \ .replace( "CONFIG_TREE_DATA", json.dumps([data])) \ .replace( "CONFIG_ZOOM", str(zoom)) \ .replace( "CONFIG_STRAIGHT_LINES", str(straight_lines).lower()) \ .replace( "CONFIG_SHOW_SUBLABELS", str(show_sublabels).lower()) \ .replace( "CONFIG_TITLE", fname) \ .replace( "CONFIG_VIS_HEIGHT", str(height)) \ .replace( "CONFIG_BG_COLOR", "#111111" if dark else "#FFFFFF") \ .replace( "CONFIG_TEXT_COLOR", '#FFFFFF' if dark else '#000000') \ .replace( "CONFIG_TEXT_RECT_COLOR", "rgba(17,17,17,0.8)" if dark else "rgba(255,255,255,1)") \ .replace( "CONFIG_MARGIN_TOP", str(margin_top)) \ .replace( "CONFIG_ROOT_Y", str(root_y)) \ .replace( "CONFIG_COLORMAP", f'''<img src="{colormap}" style=" position: absolute; top: 40px; left: 80px; height: 250px; border: 4px solid #ccc;">''' if isinstance(colormap, str) and os.path.exists(colormap) else '' ) os.makedirs(out_dir, exist_ok=True) path_html = f'{out_dir}/{fname}.html' with open(path_html, 'w') as f: f.write(html) Colors.green('==> Wrote HTML to {}'.format(path_html))
def main(): maybe_install_wordnet() datasets = data.cifar.names + data.imagenet.names + data.custom.names parser = argparse.ArgumentParser(description="PyTorch CIFAR Training") parser.add_argument("--batch-size", default=512, type=int, help="Batch size used for training") parser.add_argument( "--epochs", "-e", default=200, type=int, help="By default, lr schedule is scaled accordingly", ) parser.add_argument("--dataset", default="CIFAR10", choices=datasets) parser.add_argument("--arch", default="ResNet18", choices=list(models.get_model_choices())) parser.add_argument("--lr", default=0.1, type=float, help="learning rate") parser.add_argument("--resume", "-r", action="store_true", help="resume from checkpoint") # extra general options for main script parser.add_argument("--path-resume", default="", help="Overrides checkpoint path generation") parser.add_argument( "--name", default="", help="Name of experiment. Used for checkpoint filename") parser.add_argument( "--pretrained", action="store_true", help="Download pretrained model. Not all models support this.", ) parser.add_argument("--eval", help="eval only", action="store_true") parser.add_argument( "--dataset-test", choices=datasets, help="If not set, automatically set to train dataset", ) parser.add_argument( "--disable-test-eval", help="Allows you to run model inference on a test dataset " " different from train dataset. Use an anlayzer to define " "a metric.", action="store_true", ) # options specific to this project and its dataloaders parser.add_argument("--loss", choices=loss.names, default=["CrossEntropyLoss"], nargs="+") parser.add_argument("--metric", choices=metrics.names, default="top1") parser.add_argument("--analysis", choices=analysis.names, help="Run analysis after each epoch") # other dataset, loss or analysis specific options data.custom.add_arguments(parser) T.add_arguments(parser) loss.add_arguments(parser) analysis.add_arguments(parser) args = parser.parse_args() loss.set_default_values(args) device = "cuda" if torch.cuda.is_available() else "cpu" best_acc = 0 # best test accuracy start_epoch = 0 # start from epoch 0 or last checkpoint epoch # Data print("==> Preparing data..") dataset_train = getattr(data, args.dataset) dataset_test = getattr(data, args.dataset_test or args.dataset) transform_train = dataset_train.transform_train() transform_test = dataset_test.transform_val() dataset_train_kwargs = generate_kwargs( args, dataset_train, name=f"Dataset {dataset_train.__class__.__name__}", globals=locals(), ) dataset_test_kwargs = generate_kwargs( args, dataset_test, name=f"Dataset {dataset_test.__class__.__name__}", globals=locals(), ) trainset = dataset_train( **dataset_train_kwargs, root="./data", train=True, download=True, transform=transform_train, ) testset = dataset_test( **dataset_test_kwargs, root="./data", train=False, download=True, transform=transform_test, ) assert trainset.classes == testset.classes or args.disable_test_eval, ( trainset.classes, testset.classes, ) trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=2) testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2) Colors.cyan( f"Training with dataset {args.dataset} and {len(trainset.classes)} classes" ) Colors.cyan( f"Testing with dataset {args.dataset_test or args.dataset} and {len(testset.classes)} classes" ) # Model print("==> Building model..") model = getattr(models, args.arch) if args.pretrained: print("==> Loading pretrained model..") model = make_kwarg_optional(model, dataset=args.dataset) net = model(pretrained=True, num_classes=len(trainset.classes)) else: net = model(num_classes=len(trainset.classes)) net = net.to(device) if device == "cuda": net = torch.nn.DataParallel(net) cudnn.benchmark = True checkpoint_fname = generate_checkpoint_fname(**vars(args)) checkpoint_path = "./checkpoint/{}.pth".format(checkpoint_fname) print(f"==> Checkpoints will be saved to: {checkpoint_path}") resume_path = args.path_resume or checkpoint_path if args.resume: # Load checkpoint. print("==> Resuming from checkpoint..") assert os.path.isdir( "checkpoint"), "Error: no checkpoint directory found!" if not os.path.exists(resume_path): print("==> No checkpoint found. Skipping...") else: checkpoint = torch.load(resume_path, map_location=torch.device(device)) if "net" in checkpoint: load_state_dict(net, checkpoint["net"]) best_acc = checkpoint["acc"] start_epoch = checkpoint["epoch"] Colors.cyan( f"==> Checkpoint found for epoch {start_epoch} with accuracy " f"{best_acc} at {resume_path}") else: load_state_dict(net, checkpoint) Colors.cyan(f"==> Checkpoint found at {resume_path}") # hierarchy tree = Tree.create_from_args(args, classes=trainset.classes) # loss criterion = None for _loss in args.loss: if criterion is None and not hasattr(nn, _loss): criterion = nn.CrossEntropyLoss() class_criterion = getattr(loss, _loss) loss_kwargs = generate_kwargs( args, class_criterion, name=f"Loss {args.loss}", globals=locals(), ) criterion = class_criterion(**loss_kwargs) optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4) scheduler = optim.lr_scheduler.MultiStepLR( optimizer, milestones=[int(3 / 7.0 * args.epochs), int(5 / 7.0 * args.epochs)]) class_analysis = getattr(analysis, args.analysis or "Noop") analyzer_kwargs = generate_kwargs( args, class_analysis, name=f"Analyzer {args.analysis}", globals=locals(), ) analyzer = class_analysis(**analyzer_kwargs) metric = getattr(metrics, args.metric)() # Training @analyzer.train_function def train(epoch): if hasattr(criterion, "set_epoch"): criterion.set_epoch(epoch, args.epochs) print("\nEpoch: %d / LR: %.04f" % (epoch, scheduler.get_last_lr()[0])) net.train() train_loss = 0 metric.clear() for batch_idx, (inputs, targets) in enumerate(trainloader): inputs, targets = inputs.to(device), targets.to(device) optimizer.zero_grad() outputs = net(inputs) loss = criterion(outputs, targets) loss.backward() optimizer.step() train_loss += loss.item() metric.forward(outputs, targets) transform = trainset.transform_val_inverse().to(device) stat = analyzer.update_batch(outputs, targets, transform(inputs)) progress_bar( batch_idx, len(trainloader), "Loss: %.3f | Acc: %.3f%% (%d/%d) %s" % ( train_loss / (batch_idx + 1), 100.0 * metric.report(), metric.correct, metric.total, f"| {analyzer.name}: {stat}" if stat else "", ), ) scheduler.step() @analyzer.test_function def test(epoch, checkpoint=True): nonlocal best_acc net.eval() test_loss = 0 metric.clear() with torch.no_grad(): for batch_idx, (inputs, targets) in enumerate(testloader): inputs, targets = inputs.to(device), targets.to(device) outputs = net(inputs) if not args.disable_test_eval: loss = criterion(outputs, targets) test_loss += loss.item() metric.forward(outputs, targets) transform = testset.transform_val_inverse().to(device) stat = analyzer.update_batch(outputs, targets, transform(inputs)) progress_bar( batch_idx, len(testloader), "Loss: %.3f | Acc: %.3f%% (%d/%d) %s" % ( test_loss / (batch_idx + 1), 100.0 * metric.report(), metric.correct, metric.total, f"| {analyzer.name}: {stat}" if stat else "", ), ) # Save checkpoint. acc = 100.0 * metric.report() print("Accuracy: {}, {}/{} | Best Accurracy: {}".format( acc, metric.correct, metric.total, best_acc)) if acc > best_acc and checkpoint: Colors.green(f"Saving to {checkpoint_fname} ({acc})..") state = { "net": net.state_dict(), "acc": acc, "epoch": epoch, } os.makedirs("checkpoint", exist_ok=True) torch.save(state, f"./checkpoint/{checkpoint_fname}.pth") best_acc = acc if args.disable_test_eval and (not args.analysis or args.analysis == "Noop"): Colors.red( " * Warning: `disable_test_eval` is used but no custom metric " "`--analysis` is supplied. I suggest supplying an analysis to perform " " custom loss and accuracy calculation.") if args.eval: if not args.resume and not args.pretrained: Colors.red(" * Warning: Model is not loaded from checkpoint. " "Use --resume or --pretrained (if supported)") with analyzer.epoch_context(0): test(0, checkpoint=False) else: for epoch in range(start_epoch, args.epochs): with analyzer.epoch_context(epoch): train(epoch) test(epoch) print(f"Best accuracy: {best_acc} // Checkpoint name: {checkpoint_fname}")