def main(_run, _log): args = argparse.Namespace(**_run.config) args = post_config_hook(args, _run) args.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") root = "./datasets" train_sampler = None if args.dataset == "STL10": train_dataset = torchvision.datasets.STL10( root, split="unlabeled", download=True, transform=TransformsSimCLR() ) elif args.dataset == "CIFAR10": train_dataset = torchvision.datasets.CIFAR10( root, download=True, transform=TransformsSimCLR() ) else: raise NotImplementedError train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), drop_last=True, num_workers=args.workers, sampler=train_sampler, ) model, optimizer, scheduler = load_model(args, train_loader) tb_dir = os.path.join(args.out_dir, _run.experiment_info["name"]) os.makedirs(tb_dir) writer = SummaryWriter(log_dir=tb_dir) mask = mask_correlated_samples(args) criterion = NT_Xent(args.batch_size, args.temperature, mask, args.device) args.global_step = 0 args.current_epoch = 0 for epoch in range(args.start_epoch, args.epochs): lr = optimizer.param_groups[0]['lr'] loss_epoch = train(args, train_loader, model, criterion, optimizer, writer) if scheduler: scheduler.step() if epoch % 10 == 0: save_model(args, model, optimizer) writer.add_scalar("Loss/train", loss_epoch / len(train_loader), epoch) writer.add_scalar("Misc/learning_rate", lr, epoch) print( f"Epoch [{epoch}/{args.epochs}]\t Loss: {loss_epoch / len(train_loader)}\t lr: {round(lr, 5)}" ) args.current_epoch += 1 ## end training save_model(args, model, optimizer)
def main(_run, _log): args = argparse.Namespace(**_run.config) args = post_config_hook(args, _run) args.device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu") root = "./datasets" simclr_model = load_model(args, reload_model=True) simclr_model = simclr_model.to(args.device) simclr_model.eval() ## Logistic Regression n_classes = 10 # stl-10 model = LogisticRegression(simclr_model.n_features, n_classes) model = model.to(args.device) optimizer = torch.optim.Adam(model.parameters(), lr=3e-4) criterion = torch.nn.CrossEntropyLoss() train_dataset = torchvision.datasets.STL10( root, split="train", download=True, transform=torchvision.transforms.ToTensor()) test_dataset = torchvision.datasets.STL10( root, split="test", download=True, transform=torchvision.transforms.ToTensor()) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.logistic_batch_size, drop_last=True, num_workers=args.workers, ) test_loader = torch.utils.data.DataLoader( test_dataset, batch_size=args.logistic_batch_size, drop_last=True, num_workers=args.workers, ) for epoch in range(args.logistic_epochs): loss_epoch, accuracy_epoch = train(args, train_loader, simclr_model, model, criterion, optimizer) print( f"Epoch [{epoch}/{args.logistic_epochs}]\t Loss: {loss_epoch / len(train_loader)}\t Accuracy: {accuracy_epoch / len(train_loader)}" ) # final testing loss_epoch, accuracy_epoch = test(args, test_loader, simclr_model, model, criterion, optimizer) print( f"[FINAL]\t Loss: {loss_epoch / len(test_loader)}\t Accuracy: {accuracy_epoch / len(test_loader)}" )
def main(_run, _log): args = argparse.Namespace(**_run.config) args = post_config_hook(args, _run) args.device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu") root = "./datasets" model = load_model(args) model = model.to(args.device) optimizer = torch.optim.Adam(model.parameters(), lr=3e-4) # TODO: LARS train_sampler = None train_dataset = torchvision.datasets.STL10(root, split="unlabeled", download=True, transform=TransformsSimCLR()) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), drop_last=True, num_workers=args.workers, sampler=train_sampler, ) tb_dir = os.path.join(args.out_dir, _run.experiment_info["name"]) os.makedirs(tb_dir) writer = SummaryWriter(log_dir=tb_dir) mask = mask_correlated_samples(args) criterion = NT_Xent(args.batch_size, args.temperature, mask, args.device) args.global_step = 0 args.current_epoch = 0 for epoch in range(args.start_epoch, args.epochs): loss_epoch = train(args, train_loader, model, criterion, optimizer, writer) writer.add_scalar("Loss/train", loss_epoch / len(train_loader), epoch) if epoch % 10 == 0: save_model(args, model, optimizer) print( f"Epoch [{epoch}/{args.epochs}]\t Loss: {loss_epoch / len(train_loader)}" ) args.current_epoch += 1 ## end training save_model(args, model, optimizer)
def main(_run, _log): args = argparse.Namespace(**_run.config) args = post_config_hook(args, _run) args.device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu") root = "./datasets" train_sampler = None valid_sampler = None if args.dataset == "STL10": dataset = torchvision.datasets.STL10( root, split="train", download=True, transform=TransformsSimCLR(size=224).test_transform, ) test_dataset = torchvision.datasets.STL10( root, split="test", download=True, transform=TransformsSimCLR(size=224).test_transform, ) elif args.dataset == "CIFAR10": dataset = torchvision.datasets.CIFAR10( root, train=True, download=True, transform=TransformsSimCLR(size=224).test_transform, ) test_dataset = torchvision.datasets.CIFAR10( root, train=False, download=True, transform=TransformsSimCLR(size=224).test_transform, ) elif args.dataset == "MATEK": dataset, train_sampler, valid_sampler = MatekDataset( root=root, transforms=TransformsSimCLR(size=128).test_transform, test_size=args.test_size).get_dataset() elif args.dataset == "JURKAT": dataset, train_sampler, valid_sampler = JurkatDataset( root=root, transforms=TransformsSimCLR(size=64).test_transform, test_size=args.test_size).get_dataset() elif args.dataset == "PLASMODIUM": dataset, train_sampler, valid_sampler = PlasmodiumDataset( root=root, transforms=TransformsSimCLR(size=128).test_transform, test_size=args.test_size).get_dataset() else: raise NotImplementedError train_loader = torch.utils.data.DataLoader( dataset, batch_size=args.logistic_batch_size, shuffle=(train_sampler is None), drop_last=True, num_workers=args.workers, sampler=train_sampler, ) test_loader = torch.utils.data.DataLoader( dataset, batch_size=args.logistic_batch_size, shuffle=(valid_sampler is None), drop_last=True, num_workers=args.workers, sampler=valid_sampler, ) simclr_model, _, _ = load_model(args, train_loader, reload_model=True) simclr_model = simclr_model.to(args.device) simclr_model.eval() # Logistic Regression n_classes = args.n_classes # stl-10 model = LogisticRegression(simclr_model.n_features, n_classes) model = model.to(args.device) optimizer = torch.optim.Adam(model.parameters(), lr=3e-4) criterion = torch.nn.CrossEntropyLoss() print("### Creating features from pre-trained context model ###") (train_X, train_y, test_X, test_y) = get_features(simclr_model, train_loader, test_loader, args.device) arr_train_loader, arr_test_loader = create_data_loaders_from_arrays( train_X, train_y, test_X, test_y, args.logistic_batch_size) for epoch in range(args.logistic_epochs): loss_epoch, accuracy_epoch = train(args, arr_train_loader, simclr_model, model, criterion, optimizer) print( f"Epoch [{epoch}/{args.logistic_epochs}]\t Loss: {loss_epoch / len(train_loader)}\t Accuracy: {accuracy_epoch / len(train_loader)}" ) # final testing loss_epoch, accuracy_epoch, report = test(args, arr_test_loader, simclr_model, model, criterion, optimizer) print( f"[FINAL]\t Loss: {loss_epoch / len(test_loader)}\t Accuracy: {accuracy_epoch / len(test_loader)}" ) print(report)
def main(_run, _log): args = argparse.Namespace(**_run.config) args = post_config_hook(args, _run) args.device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu") args.n_gpu = torch.cuda.device_count() train_sampler = None if args.dataset == "STL10": train_dataset = torchvision.datasets.STL10( root=args.dataset_root, split="unlabeled", download=True, transform=TransformsSimCLR(size=96)) elif args.dataset == "CIFAR10": train_dataset = torchvision.datasets.CIFAR10( root=args.dataset_root, download=True, transform=TransformsSimCLR(size=32)) elif args.dataset == "MATEK": train_dataset, _ = MatekDataset( root=args.dataset_root, transforms=TransformsSimCLR(size=128)).get_dataset() elif args.dataset == "JURKAT": train_dataset, _ = JurkatDataset( root=args.dataset_root, transforms=TransformsSimCLR(size=64)).get_dataset() elif args.dataset == "PLASMODIUM": train_dataset, _ = PlasmodiumDataset( root=args.dataset_root, transforms=TransformsSimCLR(size=128)).get_dataset() else: raise NotImplementedError train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), drop_last=True, num_workers=args.workers, sampler=train_sampler, ) model, optimizer, scheduler = load_model(args, train_loader) print(f"Using {args.n_gpu}'s") if args.n_gpu > 1: model = torch.nn.DataParallel(model) model = convert_model(model) model = model.to(args.device) print(model) tb_dir = os.path.join(args.out_dir, _run.experiment_info["name"]) os.makedirs(tb_dir) writer = SummaryWriter(log_dir=tb_dir) criterion = NT_Xent(args.batch_size, args.temperature, args.device) args.global_step = 0 args.current_epoch = 0 for epoch in range(args.start_epoch, args.epochs): lr = optimizer.param_groups[0]['lr'] loss_epoch = train(args, train_loader, model, criterion, optimizer, writer) if scheduler: scheduler.step() if epoch % 10 == 0: save_model(args, model, optimizer) writer.add_scalar("Loss/train", loss_epoch / len(train_loader), epoch) writer.add_scalar("Misc/learning_rate", lr, epoch) print( f"Epoch [{epoch}/{args.epochs}]\t Loss: {loss_epoch / len(train_loader)}\t lr: {round(lr, 5)}" ) args.current_epoch += 1 save_model(args, model, optimizer)