def main(cfg: OmegaConf): logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) stream_handler = logging.StreamHandler() stream_handler.setLevel(logging.INFO) stream_handler.terminator = "" logger.addHandler(stream_handler) init_ddp(cfg) check_hydra_conf(cfg) seed = cfg["parameter"]["seed"] np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False rank = cfg["distributed"]["local_rank"] logger.info("Using {}".format(rank)) transform = SimCLRTransforms(strength=cfg["experiment"]["strength"]) root = "~/pytorch_datasets" if cfg["experiment"]["name"] == "cifar10": training_dataset = torchvision.datasets.CIFAR10(root=root, train=True, download=True, transform=transform) elif cfg["experiment"]["name"] == "cifar100": training_dataset = torchvision.datasets.CIFAR100(root=root, train=True, download=True, transform=transform) else: assert cfg["experiment"]["name"] in {"cifar10", "cifar100"} sampler = torch.utils.data.distributed.DistributedSampler(training_dataset, shuffle=True) training_data_loader = DataLoader( dataset=training_dataset, sampler=sampler, num_workers=cfg["parameter"]["num_workers"], batch_size=cfg["experiment"]["batches"], pin_memory=True, drop_last=True, ) model = ContrastiveModel(base_cnn=cfg["experiment"]["base_cnn"], d=cfg["parameter"]["d"]) model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) model = model.to(rank) model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank]) train(cfg, training_data_loader, model)
def main(cfg: OmegaConf): logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) stream_handler = logging.StreamHandler() stream_handler.setLevel(logging.INFO) stream_handler.terminator = "" logger.addHandler(stream_handler) check_hydra_conf(cfg) seed = cfg["parameter"]["seed"] torch.manual_seed(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False use_cuda = cfg["parameter"]["use_cuda"] and torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") logger.info("Using {}".format(device)) transform = transforms.Compose([ transforms.ToTensor(), ]) root = "~/pytorch_datasets" if cfg["experiment"]["name"] == "cifar10": training_dataset = torchvision.datasets.CIFAR10(root=root, train=True, download=True, transform=transform) val_dataset = torchvision.datasets.CIFAR10(root=root, train=False, download=True, transform=transform) num_classes = 10 elif cfg["experiment"]["name"] == "cifar100": training_dataset = torchvision.datasets.CIFAR100(root=root, train=True, download=True, transform=transform) val_dataset = torchvision.datasets.CIFAR100(root=root, train=False, download=True, transform=transform) num_classes = 100 else: assert cfg["experiment"]["name"] in {"cifar10", "cifar100"} training_data_loader = DataLoader( dataset=training_dataset, batch_size=cfg["experiment"]["batches"], shuffle=True, ) val_data_loader = DataLoader( dataset=val_dataset, batch_size=cfg["experiment"]["batches"], shuffle=False, ) classification_results = {} top_k = cfg["parameter"]["top_k"] for weights_path in Path(cfg["experiment"]["target_dir"]).glob("*.pt"): key = str(weights_path).split("/")[-1] logger.info("Evaluation by using {}".format(key)) model = ContrastiveModel(base_cnn=cfg["experiment"]["base_cnn"], d=cfg["parameter"]["d"]) model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) model = model.to(device) state_dict = torch.load(weights_path) state_dict = { k.replace("module.", ""): v for k, v in state_dict.items() } # load weights trained on self-supervised task if use_cuda: model.load_state_dict(state_dict, strict=False) else: model.load_state_dict(state_dict, strict=False, map_location=device) downstream_training_dataset = DownstreamDataset( *convert_vectors(cfg, training_data_loader, model, device)) downstream_val_dataset = DownstreamDataset( *convert_vectors(cfg, val_data_loader, model, device)) downstream_training_data_loader = DataLoader( dataset=downstream_training_dataset, batch_size=cfg["experiment"]["batches"], shuffle=True, ) downstream_val_data_loader = DataLoader( dataset=downstream_val_dataset, batch_size=cfg["experiment"]["batches"], shuffle=False, ) if cfg["parameter"]["classifier"] == "centroid": classifier = CentroidClassifier( weights=CentroidClassifier.create_weights( downstream_training_dataset, num_classes=num_classes).to( device)) train_acc, train_top_k_acc = centroid_eval( downstream_training_data_loader, device, classifier, top_k) val_acc, val_top_k_acc = centroid_eval(downstream_val_data_loader, device, classifier, top_k) classification_results[key] = { "train_acc": train_acc, "train_top_{}_acc".format(top_k): train_top_k_acc, "val_acc": val_acc, "val_top_{}_acc".format(top_k): val_top_k_acc } logger.info("train acc: {}, val acc: {}".format( train_acc, val_acc)) else: if cfg["parameter"]["use_full_encoder"]: num_last_units = model.g[-1].out_features else: num_last_units = model.g[0].in_features if cfg["parameter"]["classifier"] == "linear": classifier = LinearClassifier(num_last_units, num_classes).to(device) elif cfg["parameter"]["classifier"].replace("-", "") == "nonlinear": classifier = NonLinearClassifier(num_last_units, num_classes).to(device) train_accuracies, train_top_k_accuracies, train_losses, val_accuracies, val_top_k_accuracies, val_losses = \ learnable_eval(cfg, classifier, downstream_training_data_loader, downstream_val_data_loader, device) classification_results[key] = { "train_accuracies": train_accuracies, "val_accuracies": val_accuracies, "train_losses": train_losses, "val_losses": val_losses, "train_top_{}_accuracies".format(top_k): train_top_k_accuracies, "val_top_{}_accuracies".format(top_k): val_top_k_accuracies, "lowest_val_loss": min(val_losses), "highest_val_acc": max(val_accuracies), "highest_val_top_k_acc": max(val_top_k_accuracies) } logger.info("train acc: {}, val acc: {}".format( max(train_accuracies), max(val_accuracies))) fname = cfg["parameter"]["classification_results_json_fname"] with open(fname, "w") as f: json.dump(classification_results, f)