def init_data(self): print('Initialize dataset...') self.train_transform = data.get_transform(self.args.image_size, self.args.train_transform) self.test_transform = data.get_transform(self.args.image_size, self.args.test_transform) # load base dataset self.base_dataset, self.test_dataset = data.load_base_dataset( self.args) self.base_dataset.transform = self.train_transform self.test_dataset.transform = self.test_transform # split to train/val/pool set if self.args.init_size is None: self.train_idx = list(range(len(self.base_dataset))) self.val_idx = [] self.pool_idx = [] self.args.init_size = len(self.base_dataset) self.args.per_size = 0 self.args.max_size = len(self.base_dataset) else: self.train_idx, self.val_idx, self.pool_idx = data.split_dataset( self.base_dataset, self.args.ny, self.args.init_size, self.args.val_size) if self.args.max_size is None: self.args.per_size = 0 self.args.max_size = self.args.init_size # define trainset and pool self.trainset = data_utils.Subset(self.base_dataset, self.train_idx) self.valset = data_utils.Subset(self.base_dataset, self.val_idx) self.pool = data_utils.Subset(self.base_dataset, self.pool_idx)
def __init__(self, config, split, model_name='bert', pretrained='bert-base-uncased', finetuned=None): super(HuggingFaceTransformerExtractor, self).__init__(config, split, bs=5, collate_fn=TextCollator()) self.pretrained = pretrained self.finetuned = finetuned self.model_name = model_name self.config = config roots, ids = data.get_paths(config) data_name = config['dataset']['name'] transform = data.get_transform(data_name, 'val', config) collate_fn = data.Collate(config) self.loader = data.get_loader_single(data_name, split, roots[split]['img'], roots[split]['cap'], transform, ids=ids[split], batch_size=32, shuffle=False, num_workers=4, collate_fn=collate_fn)
def __init__(self): self.gan = Adaptor() self.dataloader = data.IconDataset("./datasets", batchSize=12) self.transform = data.get_transform() self.edge_transform = data.get_edge_transform() self.input = self.dataloader[0]
def prep_data(): print('Datasets preparation...') train_transform, valid_transform = get_transform( train=True), get_transform(train=False) dataset = DbdImageDataset('data/', train_transform) dataset_test = DbdImageDataset('data/val', valid_transform) data_loader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=4, collate_fn=collate_fn) data_loader_test = DataLoader(dataset_test, batch_size=1, shuffle=False, num_workers=4, collate_fn=collate_fn) return data_loader, data_loader_test
def optimize(opt): dataset_name = 'cat' generator_name = 'stylegan2' transform = data.get_transform(dataset_name, 'im2tensor') dset = data.get_dataset(dataset_name, opt.partition, load_w=False, transform=transform) total = len(dset) if opt.indices is None: start_idx = 0 end_idx = total else: start_idx = opt.indices[0] end_idx = opt.indices[1] print("Optimizing dataset partition %s items %d to %d" % (opt.partition, start_idx, end_idx)) generator = domain_generator.define_generator(generator_name, dataset_name, load_encoder=True) util.set_requires_grad(False, generator.generator) util.set_requires_grad(False, generator.encoder) for i in range(start_idx, end_idx): (im, label, path) = dset[i] img_filename = os.path.splitext(os.path.basename(path))[0] print("Running %d / %d images: %s" % (i, end_idx, img_filename)) output_filename = os.path.join(opt.w_path, img_filename) if os.path.isfile(output_filename + '.pth'): print(output_filename + '.pth found... skipping') continue # cat face dataset is already centered centered_im = im[None].cuda() # find zero values to estimate the mask mask = torch.ones_like(centered_im) mask[torch.where( torch.sum(torch.abs(centered_im), axis=0, keepdims=True) < 0.02 )] = 0 mask = mask[:, :1, :, :].cuda() ckpt, loss = generator.optimize(centered_im, mask=mask) w_optimized = ckpt['current_z'] loss = np.array(loss).squeeze() im_optimized = renormalize.as_image(ckpt['current_x'][0]) torch.save({'w': w_optimized.detach().cpu()}, output_filename + '.pth') np.savez(output_filename + '_loss.npz', loss=loss) im_optimized.save(output_filename + '_optimized_im.png')
def inference(model, image): model.eval() transform = data.get_transform('test') x = transform(image) x = x.unsqueeze(0) x = x.to(model.device) output = model(x) predict = torch.argmax(output, -1).item() return predict
def optimize(opt): dataset_name = 'celebahq' generator_name = 'stylegan2' transform = data.get_transform(dataset_name, 'imval') # we don't need the labels, so attribute doesn't really matter here dset = data.get_dataset(dataset_name, opt.partition, 'Smiling', load_w=False, return_path=True, transform=transform) total = len(dset) if opt.indices is None: start_idx = 0 end_idx = total else: start_idx = opt.indices[0] end_idx = opt.indices[1] print("Optimizing dataset partition %s items %d to %d" % (opt.partition, start_idx, end_idx)) generator = domain_generator.define_generator(generator_name, dataset_name, load_encoder=True) util.set_requires_grad(False, generator.generator) util.set_requires_grad(False, generator.encoder) for i in range(start_idx, end_idx): (image, label, path) = dset[i] image = image[None].cuda() img_filename = os.path.splitext(os.path.basename(path))[0] print("Running %d / %d images: %s" % (i, end_idx, img_filename)) output_filename = os.path.join(opt.w_path, img_filename) if os.path.isfile(output_filename + '.pth'): print(output_filename + '.pth found... skipping') continue ckpt, loss = generator.optimize(image, mask=None) w_optimized = ckpt['current_z'] loss = np.array(loss).squeeze() im_optimized = renormalize.as_image(ckpt['current_x'][0]) torch.save({'w': w_optimized.detach().cpu()}, output_filename + '.pth') np.savez(output_filename + '_loss.npz', loss=loss) im_optimized.save(output_filename + '_optimized_im.png')
def do_predict(pil_image, model, device, top_k=1): # def validate_epoch(loader, model, crit, device, desc='Validating'): model.train(False) model.eval() preprocess = data.get_transform(arch=model.arch, training=False) image = preprocess(pil_image) image.unsqueeze_(0) # adds a dimension for batch, else it fails image = image.to(device) output = model(image) scores, labels = torch.topk(output.data, top_k) probs = F.softmax(scores.data, dim=1) probs = probs.to('cpu')[0].numpy() labels = labels.to('cpu')[0].numpy().astype(str) return labels, probs
def __init__(self, context: det.TrialContext) -> None: self.context = context # Create a unique download directory for each rank so they don't # overwrite each other. self.download_directory = f"/tmp/data-rank{self.context.distributed.get_rank()}" download_data( download_directory=self.download_directory, data_config=self.context.get_data_config(), ) dataset = PennFudanDataset(self.download_directory + "/PennFudanPed", get_transform()) # Split 80/20 into training and validation datasets. train_size = int(0.8 * len(dataset)) test_size = len(dataset) - train_size self.dataset_train, self.dataset_val = torch.utils.data.random_split( dataset, [train_size, test_size] )
def optimize(opt): dataset_name = 'cifar10' generator_name = 'stylegan2-cc' # class conditional stylegan transform = data.get_transform(dataset_name, 'imval') dset = data.get_dataset(dataset_name, opt.partition, load_w=False, transform=transform) total = len(dset) if opt.indices is None: start_idx = 0 end_idx = total else: start_idx = opt.indices[0] end_idx = opt.indices[1] generator = domain_generator.define_generator(generator_name, dataset_name, load_encoder=False) util.set_requires_grad(False, generator.generator) resnet = domain_classifier.define_classifier(dataset_name, 'imageclassifier') ### iterate ### for i in range(start_idx, end_idx): img, label = dset[i] print("Running img %d/%d" % (i, len(dset))) filename = os.path.join(opt.w_path, '%s_%06d.npy' % (opt.partition, i)) if os.path.isfile(filename): print(filename + ' found... skipping') continue img = img[None].cuda() with torch.no_grad(): pred_logit = resnet(img) _, pred_label = pred_logit.max(1) pred_label = pred_label.item() print("True label %d prd label %d" % (label, pred_label)) ckpt, loss = generator.optimize(img, pred_label) current_z = ckpt['current_z'].detach().cpu().numpy() np.save(filename, current_z)
def __init__(self, context: PyTorchTrialContext) -> None: self.context = context # Create a unique download directory for each rank so they don't # overwrite each other. self.download_directory = f"/tmp/data-rank{self.context.distributed.get_rank()}" download_data( download_directory=self.download_directory, data_config=self.context.get_data_config(), ) dataset = PennFudanDataset(self.download_directory + "/PennFudanPed", get_transform()) # Split 80/20 into training and validation datasets. train_size = int(0.8 * len(dataset)) test_size = len(dataset) - train_size self.dataset_train, self.dataset_val = torch.utils.data.random_split( dataset, [train_size, test_size] ) model = fasterrcnn_resnet50_fpn(pretrained=True) # Replace the classifier with a new two-class classifier. There are # only two "classes": pedestrian and background. num_classes = 2 in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) # Wrap the model. self.model = self.context.wrap_model(model) # Wrap the optimizer. self.optimizer = self.context.wrap_optimizer(torch.optim.SGD( self.model.parameters(), lr=self.context.get_hparam("learning_rate"), momentum=self.context.get_hparam("momentum"), weight_decay=self.context.get_hparam("weight_decay"), )) # Wrap the LR scheduler. self.lr_scheduler = self.context.wrap_lr_scheduler( torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=3, gamma=0.1), step_mode=LRScheduler.StepMode.STEP_EVERY_EPOCH )
###################################################################### # Formatting the Data # ------------------- # # This is a good place to apply transformations to the data. For the # waveform, we downsample the audio for faster processing without losing # too much of the classification power. # # We don’t need to apply other transformations here. It is common for some # datasets though to have to reduce the number of channels (say from # stereo to mono) by either taking the mean along the channel dimension, # or simply keeping only one of the channels. Since SpeechCommands uses a # single channel for audio, this is not needed here. # transform = data.get_transform(sample_rate) transformed = transform(waveform) # ipd.Audio(transformed.numpy(), rate=new_sample_rate) ###################################################################### # We are encoding each word using its index in the list of labels. # word_start = "yes" index = data.label_to_index(word_start, labels) word_recovered = data.index_to_label(index, labels) print(word_start, "-->", index, "-->", word_recovered) ######################################################################
def main(args): # use_cuda = not args.no_cuda and torch.cuda.is_available() set_random_seed(args.seed) device = torch.device("cuda" if use_cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {} if args.dataset == 'mnist': train_data = get_dataset('mnist-train', args.dataroot) test_data = get_dataset('mnist-test', args.dataroot) train_tr = test_tr = get_transform('mnist_normalize') if args.dataset == 'cifar10': train_tr_name = 'cifar_augment_normalize' if args.data_augmentation else 'cifar_normalize' train_data = get_dataset('cifar10-train', args.dataroot) test_data = get_dataset('cifar10-test', args.dataroot) train_tr = get_transform(train_tr_name) test_tr = get_transform('cifar_normalize') if args.dataset == 'cifar-fs-train': train_tr_name = 'cifar_augment_normalize' if args.data_augmentation else 'cifar_normalize' train_data = get_dataset('cifar-fs-train-train', args.dataroot) test_data = get_dataset('cifar-fs-train-test', args.dataroot) train_tr = get_transform(train_tr_name) test_tr = get_transform('cifar_normalize') if args.dataset == 'miniimagenet': train_data = get_dataset('miniimagenet-train-train', args.dataroot) test_data = get_dataset('miniimagenet-train-test', args.dataroot) train_tr = get_transform('cifar_augment_normalize_84' if args.data_augmentation else 'cifar_normalize') test_tr = get_transform('cifar_normalize') model = ResNetClassifier(train_data['n_classes'], train_data['im_size']).to(device) if args.ckpt_path != '': loaded = torch.load(args.ckpt_path) model.load_state_dict(loaded) ipdb.set_trace() if args.eval: acc = test(args, model, device, test_loader, args.n_eval_batches) print("Eval Acc: ", acc) sys.exit() # Trace logging mkdir(args.output_dir) eval_fieldnames = ['global_iteration','val_acc','train_acc'] eval_logger = CSVLogger(every=1, fieldnames=eval_fieldnames, resume=args.resume, filename=os.path.join(args.output_dir, 'eval_log.csv')) wandb.run.name = os.path.basename(args.output_dir) wandb.run.save() wandb.watch(model) if args.optim == 'adadelta': optimizer = optim.Adadelta(model.parameters(), lr=args.lr) elif args.optim == 'adam': optimizer = optim.Adam(model.parameters(), lr=args.lr) elif args.optim == 'sgd': optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, nesterov=True, weight_decay=5e-4) if args.dataset == 'mnist': scheduler = StepLR(optimizer, step_size=1, gamma=.7) else: scheduler = MultiStepLR(optimizer, milestones=[60, 120, 160], gamma=0.2) start_epoch = 1 if args.resume: last_ckpt_path = os.path.join(args.output_dir, 'last_ckpt.pt') if os.path.exists(last_ckpt_path): loaded = torch.load(last_ckpt_path) model.load_state_dict(loaded['model_sd']) optimizer.load_state_dict(loaded['optimizer_sd']) scheduler.load_state_dict(loaded['scheduler_sd']) start_epoch = loaded['epoch'] # It's important to set seed again before training b/c dataloading code # might have reset the seed. set_random_seed(args.seed) best_val = 0 if args.db: scheduler = MultiStepLR(optimizer, milestones=[1, 2, 3, 4], gamma=0.1) args.epochs = 5 for epoch in range(start_epoch, args.epochs + 1): if epoch % args.ckpt_every == 0: torch.save(model.state_dict(), os.path.join(args.output_dir , f"ckpt_{epoch}.pt")) stats_dict = {'global_iteration':epoch} val = stats_dict['val_acc'] = test(args, model, device, test_data, test_tr, args.n_eval_batches) stats_dict['train_acc'] = test(args, model, device, train_data, test_tr, args.n_eval_batches) grid = make_grid(torch.stack([train_tr(x) for x in train_data['x'][:30]]), nrow=6).permute(1,2,0).numpy() img_dict = {"examples": [wandb.Image(grid, caption="Data batch")]} wandb.log(stats_dict) wandb.log(img_dict) eval_logger.writerow(stats_dict) plot_csv(eval_logger.filename, os.path.join(args.output_dir, 'iteration_plots.png')) train(args, model, device, train_data, train_tr, optimizer, epoch) scheduler.step(epoch) if val > best_val: best_val = val torch.save(model.state_dict(), os.path.join(args.output_dir , f"ckpt_best.pt")) # For `resume` model.cpu() torch.save({ 'model_sd': model.state_dict(), 'optimizer_sd': optimizer.state_dict(), 'scheduler_sd': scheduler.state_dict(), 'epoch': epoch + 1 }, os.path.join(args.output_dir, "last_ckpt.pt")) model.to(device)
def main(): parser = argparse.ArgumentParser(description='PyTorch MNIST Example') parser.add_argument('--batch-size', type=int, default=64, metavar='N', help='input batch size for training (default: 64)') parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N', help='input batch size for testing (default: 1000)') parser.add_argument('--seed', type=int, default=1, metavar='S', help='random seed (default: 1)') parser.add_argument('--log-interval', type=int, default=10, metavar='N', help='how many batches to wait before logging training status') parser.add_argument('--save-model', action='store_true', default=False, help='For Saving the current Model') parser.add_argument('--dm_path', type=str, default='') parser.add_argument('--oec_path', type=str, default='') parser.add_argument('--episodic_ood_eval', type=int, default=0) parser.add_argument('--episodic_in_distr', type = str, default='meta-test', choices=['meta-test','meta-train']) # DM parser.add_argument('--dm_g_magnitude', type=float, default=0) parser.add_argument('--dm_ls', type=str, default='-') parser.add_argument('--db', type = int, default=0) parser.add_argument('--tag', type = str, default='') parser.add_argument('--n_episodes', type = int, default=100) parser.add_argument('--n_ways', type = int, default=5) parser.add_argument('--n_shots', type = int, default=5) # Required parser.add_argument('--dataroot', required=True) parser.add_argument('--output_dir', required=True) parser.add_argument('--dataset', required=True, choices=['mnist','cifar10', 'cifar100', 'cifar-fs', 'cifar-64', 'miniimagenet']) parser.add_argument('--ood_methods', type=str, required=True, help='comma separated list of method names e.g., `mpp,DM-all') ## Pretrained model paths parser.add_argument('--fsmodel_path', required=True) parser.add_argument('--fsmodel_name', required=True, type=str, choices=['protonet', 'maml','baseline','baseline-pn']) parser.add_argument('--classifier_path', required=True) parser.add_argument('--glow_dir', required=True) parser.add_argument('--ooe_only', type=int, default=0) args = parser.parse_args() use_cuda = True mkdir(args.output_dir) torch.manual_seed(args.seed) device = torch.device("cuda" if use_cuda else "cpu") if args.dataset == 'mnist': test_data = get_dataset('mnist-test', args.dataroot) out_list = ['gaussian', 'rademacher', 'texture3', 'svhn', 'notMNIST'] tr = get_transform('mnist_resize_normalize') if args.dataset.startswith('cifar'): out_list = ['gaussian', 'rademacher', 'texture3', 'svhn','tinyimagenet','lsun'] # out_list = ['svhn'] normalize = cifar_normalize if args.dataset == 'cifar10': train_data = get_dataset('cifar10-train', args.dataroot) test_data = get_dataset('cifar10-test', args.dataroot) if args.dataset == 'cifar100': train_data = get_dataset('cifar100-train', args.dataroot) test_data = get_dataset('cifar100-test', args.dataroot) if args.dataset == 'cifar-fs': train_data = get_dataset('cifar-fs-train-train', args.dataroot) test_data = get_dataset('cifar-fs-test', args.dataroot) if args.dataset == 'cifar-64': assert args.db train_data = get_dataset('cifar-fs-train-train', args.dataroot) test_data = get_dataset('cifar-fs-train-test', args.dataroot) tr = get_transform('cifar_resize_glow_preproc') if args.ood_methods.split(',')[0].startswith('glow') else get_transform('cifar_resize_normalize') if args.dataset == 'miniimagenet': train_data = get_dataset('miniimagenet-train-train', args.dataroot) test_data = get_dataset('miniimagenet-test', args.dataroot) out_list = ['gaussian', 'rademacher', 'texture3', 'svhn','tinyimagenet','lsun'] tr = get_transform('cifar_resize_glow_preproc') if args.ood_methods.split(',')[0].startswith('glow') else get_transform('cifar_resize_normalize_84') normalize = cifar_normalize # Models classifier = None glow = None fs_model = None ## FS Model if args.fsmodel_name in ['protonet', 'maml']: assert args.fsmodel_path != '-' fs_model = torch.load(args.fsmodel_path) encoder = fs_model.encoder ## Classifier elif args.fsmodel_name in ['baseline','baseline-pn'] : assert args.classifier_path != '-' classifier = ResNetClassifier(train_data['n_classes'], train_data['im_size']).to(device) classifier.load_state_dict(torch.load(args.classifier_path)) encoder = classifier.encoder if args.fsmodel_name == 'baseline': fs_model = BaselineFinetune(encoder, args.n_ways,args.n_shots,loss_type='dist') else: fs_model = Protonet(encoder) fs_model.to(device) fs_model.eval() args.num_feats = encoder.depth encoder.to(device) encoder.eval() if args.classifier_path != '-' and classifier is None: # for non-FS methods classifier = ResNetClassifier(train_data['n_classes'], train_data['im_size']).to(device) classifier.load_state_dict(torch.load(args.classifier_path)) if args.glow_dir != '-': # Load Glow glow_name = list(filter( lambda s: 'glow_model' in s, os.listdir(args.glow_dir)))[0] with open(os.path.join(args.glow_dir ,'hparams.json')) as json_file: hparams = json.load(json_file) # Notice Glow is 32,32,3 even for miniImageNet glow = Glow((32,32,3), hparams['hidden_channels'], hparams['K'], hparams['L'], hparams['actnorm_scale'], hparams['flow_permutation'], hparams['flow_coupling'], hparams['LU_decomposed'], train_data['n_classes'], hparams['learn_top'], hparams['y_condition']) glow.load_state_dict(torch.load(os.path.join(args.glow_dir, glow_name))) glow.set_actnorm_init() glow = glow.to(device) glow = glow.eval() # Verify Acc (just making sure models are loaded properly) if classifier is not None and not args.ood_methods.split(',')[0].startswith('glow'): preds = classifier(torch.stack([tr(x) for x in train_data['x'][:args.test_batch_size]]).to(device)).max(-1)[1] print("Train Acc: ", (preds.detach().cpu().numpy()==np.array(train_data['y'])[:args.test_batch_size]).mean()) preds = classifier(torch.stack([tr(x) for x in test_data['x'][:args.test_batch_size]]).to(device)).max(-1)[1] print("Test Acc: ", (preds.detach().cpu().numpy()==np.array(test_data['y'])[:args.test_batch_size]).mean()) # Confidence functions for OOD confidence_funcs = OrderedDict() # (name, (func, use_support, kwargs)) for ood_method in args.ood_methods.split(','): no_grad = True if ood_method.startswith('DM'): deep_mahala_obj = DeepMahala(train_data['x'], train_data['y'], tr, encoder, device,num_feats=args.num_feats, num_classes=train_data['n_classes'], pretrained_path=args.dm_path, fit=True, normalize=normalize) if ood_method.startswith('deep-ed'): no_grad=False deep_mahala_obj = DeepMahala(train_data['x'], train_data['y'], tr, encoder, device,num_feats=args.num_feats, num_classes=train_data['n_classes'], pretrained_path=args.dm_path, fit=False, normalize=normalize) if ood_method == 'MPP': confidence_funcs['MPP'] = BaseConfidence(lambda x:mpp(classifier, x)) elif ood_method == 'Ensemble-MPP': nets = [] class PModel(nn.Module): def __init__(self, logp_model): super(PModel, self).__init__() self.logp_model = logp_model def forward(self, x): return self.logp_model(x).exp() for i in range(5): _dir = os.path.dirname(args.classifier_path) _fname = os.path.basename(args.classifier_path) path = os.path.join(_dir[:-1]+f"{i}", _fname) model = ResNetClassifier(train_data['n_classes'], train_data['im_size']) model.load_state_dict(torch.load(path)) model = PModel(model) model.eval() # nets.append(model.to(device)) ensemble = Ensemble(nets) confidence_funcs['Ensemble-MPP'] = BaseConfidence(lambda x:ensemble(x).max(-1)[0]) elif ood_method == 'DM-last': confidence_funcs['DM-last'] = DMConfidence(deep_mahala_obj, {'ls':[args.num_feats - 1],'reduction':'max'}, False).to(device) elif ood_method == 'DM-all': confidence_funcs['DM-all'] = DMConfidence(deep_mahala_obj, {'ls':[i for i in range(args.num_feats)],'reduction':'max'}, False).to(device) elif ood_method == 'glow-ll': confidence_funcs['glow-ll'] = BaseConfidence(lambda x:-glow(x)[1]) elif ood_method == 'glow-lr': from test_glow_ood import ll_to_png_code_ratio confidence_funcs['glow-lr'] = BaseConfidence(lambda x:ll_to_png_code_ratio(x, glow)) elif ood_method == 'native-spp' and args.episodic_ood_eval: if args.fsmodel_name in ['maml','baseline']: no_grad=False confidence_funcs['native-spp'] = FSCConfidence(fs_model, 'spp') elif ood_method == 'native-ed' and args.episodic_ood_eval: confidence_funcs['native-ed'] = FSCConfidence(fs_model, 'ed') elif ood_method.startswith('deep-ed') and args.episodic_ood_eval: if args.dm_ls == '-': ls = range(args.num_feats) else: ls = [int(l) for l in args.dm_ls.split(',')] kwargs = { 'ls':ls, 'reduction':'max', 'g_magnitude': args.dm_g_magnitude } dm_conf = DMConfidence(deep_mahala_obj, kwargs, True, ood_method.split('-')[-1]) dm_conf.to(device) confidence_funcs[ood_method] = dm_conf elif ood_method == 'dkde' and args.episodic_ood_eval: confidence_funcs['dkde'] = DKDEConfidence(encoder) elif ood_method == 'oec' and args.episodic_ood_eval: oec_opt = json.load( open(os.path.join(os.path.dirname(args.oec_path), 'args.json'), 'r') ) init_sample = load_episode(train_data, tr, oec_opt['data.test_way'], oec_opt['data.test_shot'], oec_opt['data.test_query'], device) if oec_opt['confidence_method'] == 'oec': oec_conf = OECConfidence(None, fs_model, init_sample, oec_opt) else: oec_conf = DeepOECConfidence(None, fs_model, init_sample, oec_opt) oec_conf.load_state_dict( torch.load(args.oec_path) ) oec_conf.eval() oec_conf.to(device) confidence_funcs['oec'] = oec_conf elif ood_method == 'oec-ensemble' and args.episodic_ood_eval: # not much more effective than 'oec' oec_opt = json.load( open(os.path.join(os.path.dirname(args.oec_path), 'args.json'), 'r') ) oec_confs = [] for e in range(5): init_sample = load_episode(train_data, tr, oec_opt['data.test_way'], oec_opt['data.test_shot'], oec_opt['data.test_query'], device) if oec_opt['confidence_method'] == 'oec': oec_conf = OECConfidence(None, fs_model, init_sample, oec_opt) else: oec_conf = DeepOECConfidence(None, fs_model, init_sample, oec_opt) # Find ckpt cdir = os.path.dirname(args.oec_path)[:-1]+f"{e}" fname = list(filter(lambda s:s.endswith('conf_best.pt'), os.listdir(cdir)))[0] oec_conf.load_state_dict( torch.load(os.path.join( cdir, fname)) ) oec_conf.eval() oec_conf.to(device) oec_confs.append(oec_conf) confidence_funcs['oec'] = Ensemble(oec_confs) else: raise # ood_method not implemented, or typo in name auroc_data = defaultdict(list) auroc_95ci_data = defaultdict(list) fpr_data = defaultdict(list) fpr_95ci_data = defaultdict(list) # Classic OOD evaluation if not args.episodic_ood_eval: for out_name in out_list: ooc_config = { 'name': out_name, 'ood_scale': 1, 'n_anom': 5000, 'cuda': False } ood_tensor = load_ood_data(ooc_config) assert len(ood_tensor) <= len(test_data['x']) in_scores = defaultdict(list) out_scores = defaultdict(list) with torch.no_grad(): for i in tqdm(range(0, len(ood_tensor), args.test_batch_size)): stop = min(args.test_batch_size, len(ood_tensor[i:])) in_x = torch.stack([tr(x) for x in test_data['x'][i:i+stop]]).to(device) out_x = torch.stack([tr(x) for x in ood_tensor[i:i+stop]]).to(device) for c, f in confidence_funcs.items(): in_scores[c].append(f.score(in_x)) out_scores[c].append(f.score(out_x)) # save ood images for debugging vutils.save_image(out_x[:100], f'non-episodic-{out_name}.png' , normalize=True, nrow=10) for c in confidence_funcs: auroc = show_ood_detection_results_softmax(torch.cat(in_scores[c]).cpu().numpy(),torch.cat(out_scores[c]).cpu().numpy())[1] print(out_name, c, ': ', auroc) # auroc_data[c].append(auroc) auroc_data['dset'].append(out_name) pandas.DataFrame(auroc_data).to_csv(os.path.join(args.output_dir,f'md_auroc_{args.ood_methods}.csv')) else: cifar_meta_train_data = get_dataset('cifar-fs-train-test', args.dataroot) cifar_meta_test_data = get_dataset('cifar-fs-test', args.dataroot) # OOD Eval if args.episodic_in_distr == 'meta-train': episodic_in_data = train_data else: episodic_in_data = test_data episodic_ood = ['ooe','cifar-fs-test', 'cifar-fs-train-test'] ood_tensors = [None] + [load_ood_data({ 'name': out_name, 'ood_scale': 1, 'n_anom': 10000, }) for out_name in episodic_ood[1:] + out_list] if args.ooe_only: all_oods = [('ooe', None)] else: all_oods = zip(episodic_ood + out_list, ood_tensors) for out_name, ood_tensor in all_oods: n_query = 15 metrics_dic = defaultdict(list) for c, f in confidence_funcs.items(): metrics_dic[c] = eval_ood_aurocs( ood_tensor, episodic_in_data, tr, args.n_ways, args.n_shots, n_query, args.n_episodes, device, f, db=args.db, out_name=out_name, no_grad=no_grad ) for c in confidence_funcs: auroc = np.mean(metrics_dic[c]['aurocs']) auroc_95ci = np.std(metrics_dic[c]['aurocs']) * 1.96 / args.n_episodes auroc_data[c].append(auroc) auroc_95ci_data[c].append(auroc_95ci) print(out_name, c, 'auroc: ', auroc, ',', auroc_95ci) fpr = np.mean(metrics_dic[c]['fprs']) fpr_95ci = np.std(metrics_dic[c]['fprs']) * 1.96 / args.n_episodes fpr_data[c].append(fpr) fpr_95ci_data[c].append(fpr_95ci) print(out_name, c, 'fpr: ', fpr, ',', fpr_95ci) auroc_data['dset'].append(out_name) fpr_data['dset'].append(out_name) auroc_95ci_data['dset'].append(out_name) fpr_95ci_data['dset'].append(out_name) pandas.DataFrame(auroc_data).to_csv(os.path.join(args.output_dir,f'{args.tag}_episodic_{args.episodic_in_distr}_{args.dm_path.split(".")[0]}_{args.ood_methods}_auroc.csv')) pandas.DataFrame(fpr_data).to_csv(os.path.join(args.output_dir,f'{args.tag}_episodic_{args.episodic_in_distr}_{args.dm_path.split(".")[0]}_{args.ood_methods}_fpr.csv')) pandas.DataFrame(auroc_95ci_data).to_csv(os.path.join(args.output_dir,f'{args.tag}_episodic_{args.episodic_in_distr}_{args.dm_path.split(".")[0]}_{args.ood_methods}_auroc_95ci.csv')) pandas.DataFrame(fpr_95ci_data).to_csv(os.path.join(args.output_dir,f'{args.tag}_episodic_{args.episodic_in_distr}_{args.dm_path.split(".")[0]}_{args.ood_methods}_fpr_95ci.csv'))
parser.add_argument('--resume', '-r', action='store_true', help='resume from checkpoint') args = parser.parse_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 save_dir = 'results/classifiers/cifar10/imageclassifier' os.makedirs(save_dir, exist_ok=True) pidfile.exit_if_job_done(save_dir) # Data print('==> Preparing data..') transform_train = data.get_transform('cifar10', 'imtrain') transform_test = data.get_transform('cifar10', 'imval') trainset = data.get_dataset('cifar10', 'train', load_w=False, transform=transform_train) trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2) # modified to use a validation partition testset = data.get_dataset('cifar10', 'val', load_w=False,
def optimize(opt): dataset_name = 'car' generator_name = 'stylegan2' transform = data.get_transform(dataset_name, 'im2tensor') # loads the PIL image dset = data.get_dataset(dataset_name, opt.partition, load_w=False, transform=None) total = len(dset) if opt.indices is None: start_idx = 0 end_idx = total else: start_idx = opt.indices[0] end_idx = opt.indices[1] print("Optimizing dataset partition %s items %d to %d" % (opt.partition, start_idx, end_idx)) generator = domain_generator.define_generator(generator_name, dataset_name, load_encoder=True) util.set_requires_grad(False, generator.generator) util.set_requires_grad(False, generator.encoder) for i in range(start_idx, end_idx): (im, label, bbox, path) = dset[i] img_filename = os.path.splitext(os.path.basename(path))[0] print("Running %d / %d images: %s" % (i, end_idx, img_filename)) output_filename = os.path.join(opt.w_path, img_filename) if os.path.isfile(output_filename + '.pth'): print(output_filename + '.pth found... skipping') continue # scale image to 512 width width, height = im.size ratio = 512 / width new_width = 512 new_height = int(ratio * height) new_im = im.resize((new_width, new_height), Image.ANTIALIAS) print(im.size) print(new_im.size) bbox = [int(x * ratio) for x in bbox] # shift to center the bbox annotation cx = (bbox[2] + bbox[0]) // 2 cy = (bbox[3] + bbox[1]) // 2 print("%d --> %d" % (cx, new_width // 2)) print("%d --> %d" % (cy, new_height // 2)) offset_x = new_width // 2 - cx offset_y = new_height // 2 - cy im_tensor = transform(new_im) im_tensor, mask = data.transforms.shift_tensor(im_tensor, offset_y, offset_x) im_tensor = data.transforms.centercrop_tensor(im_tensor, 384, 512) mask = data.transforms.centercrop_tensor(mask, 384, 512) # now image size is at most 512 x 384 (could be smaller) # center car on 512x512 tensor disp_y = (512 - im_tensor.shape[1]) // 2 disp_x = (512 - im_tensor.shape[2]) // 2 centered_im = torch.ones((3, 512, 512)) * 0 centered_im[:, disp_y:disp_y + im_tensor.shape[1], disp_x:disp_x + im_tensor.shape[2]] = im_tensor centered_mask = torch.zeros_like(centered_im) centered_mask[:, disp_y:disp_y + im_tensor.shape[1], disp_x:disp_x + im_tensor.shape[2]] = mask ckpt, loss = generator.optimize(centered_im[None].cuda(), centered_mask[:1][None].cuda()) w_optimized = ckpt['current_z'] loss = np.array(loss).squeeze() im_optimized = renormalize.as_image(ckpt['current_x'][0]) torch.save({'w': w_optimized.detach().cpu()}, output_filename + '.pth') np.savez(output_filename + '_loss.npz', loss=loss) im_optimized.save(output_filename + '_optimized_im.png')
def get_dataset( name, data_root, split='train', size=224, resolution=256, num_frames=16, dataset_type='DeepfakeFrame', sampler_type='TemporalSegmentSampler', record_set_type='DeepfakeSet', segment_count=None, normalize=True, rescale=True, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], **kwargs, ): if isinstance(name, (list, tuple)): return torch.utils.data.ConcatDataset( [ get_dataset( n, data_root, split=split, size=size, resolution=resolution, num_frames=num_frames, dataset_type=dataset_type, sampler_type=sampler_type, record_set_type=record_set_type, segment_count=segment_count, normalize=normalize, rescale=rescale, mean=mean, std=std, **kwargs, ) for n in name ] ) elif name.lower() in ['all', 'full']: names = cfg.ALL_DATASETS return get_dataset( names, data_root, split=split, size=size, resolution=resolution, num_frames=num_frames, dataset_type=dataset_type, sampler_type=sampler_type, record_set_type=record_set_type, segment_count=segment_count, normalize=normalize, rescale=rescale, mean=mean, std=std, **kwargs, ) segment_count = num_frames if segment_count is None else segment_count # if dataset_type == 'DeepfakeFaceHeatvolVideo' and (name != 'DFDC'): # dataset_type = 'DeepfakeFaceVideo' metadata = cfg.get_metadata( name, split=split, dataset_type=dataset_type, record_set_type=record_set_type, data_root=data_root, ) kwargs = {**metadata, **kwargs, 'segment_count': segment_count} Dataset = getattr(data, dataset_type, 'DeepfakeFrame') RecSet = getattr(data, record_set_type, 'DeepfakeSet') Sampler = getattr(samplers, sampler_type, 'TSNFrameSampler') r_kwargs, _ = utils.split_kwargs_by_func(RecSet, kwargs) s_kwargs, _ = utils.split_kwargs_by_func(Sampler, kwargs) record_set = RecSet(**r_kwargs) sampler = Sampler(**s_kwargs) full_kwargs = { 'record_set': record_set, 'sampler': sampler, 'transform': data.get_transform( split=split, size=size, normalize=normalize, rescale=rescale, mean=mean, std=std, ), **kwargs, } dataset_kwargs, _ = utils.split_kwargs_by_func(Dataset, full_kwargs) return Dataset(**dataset_kwargs)
def train(opt): print("Random Seed: ", opt.seed) random.seed(opt.seed) torch.manual_seed(opt.seed) cudnn.benchmark = True device = 'cuda' batch_size = int(opt.batch_size) domain = opt.domain # tensorboard os.makedirs(os.path.join(opt.outf, 'runs'), exist_ok=True) writer = SummaryWriter(log_dir=os.path.join(opt.outf, 'runs')) # datasets train_transform = data.get_transform(domain, 'imtrain') train_dset = data.get_dataset(domain, 'train', load_w=False, transform=train_transform) print("Training transform:") print(train_transform) val_transform = data.get_transform(domain, 'imval') val_dset = data.get_dataset(domain, 'val', load_w=False, transform=val_transform) print("Validation transform:") print(val_transform) train_loader = DataLoader(train_dset, batch_size=opt.batch_size, shuffle=True, pin_memory=False, num_workers=opt.workers) val_loader = DataLoader(val_dset, batch_size=opt.batch_size, shuffle=False, pin_memory=False, num_workers=opt.workers) # classifier: resnet18 model net = torchvision.models.resnet18( num_classes=len(train_dset.coarse_labels)) if not opt.train_from_scratch: state_dict = torchvision.models.utils.load_state_dict_from_url( torchvision.models.resnet.model_urls['resnet18']) del state_dict['fc.weight'] del state_dict['fc.bias'] net.load_state_dict(state_dict, strict=False) net = net.to(device) # losses + optimizers + scheduler # use smaller learning rate for the feature layers if initialized # with imagenet pretrained weights criterion = nn.CrossEntropyLoss().to(device) # loss(output, target) fc_params = [k[1] for k in net.named_parameters() if k[0].startswith('fc')] feat_params = [ k[1] for k in net.named_parameters() if not k[0].startswith('fc') ] feature_backbone_lr = opt.lr if opt.train_from_scratch else 0.1 * opt.lr print("Initial learning rate for feature backbone: %0.4f" % feature_backbone_lr) print("Initial learning rate for FC layer: %0.4f" % opt.lr) optimizer = optim.Adam([{ 'params': fc_params }, { 'params': feat_params, 'lr': feature_backbone_lr }], lr=opt.lr, betas=(opt.beta1, 0.999)) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', patience=10, min_lr=1e-6, verbose=True) start_ep = 0 best_val_acc = 0.0 best_val_epoch = 0 for epoch in pbar(range(start_ep, opt.niter + 1)): # average meter for train/val loss and train/val acc metrics = dict(train_loss=util.AverageMeter(), val_loss=util.AverageMeter(), train_acc=util.AverageMeter(), val_acc=util.AverageMeter()) # train loop for step, item in enumerate(pbar(train_loader)): im = item[0].cuda() label = item[1].cuda() net.zero_grad() output = net(im) loss = criterion(output, label) accs, _ = accuracy(output, label, topk=(1, )) metrics['train_loss'].update(loss, n=len(label)) metrics['train_acc'].update(accs[0], n=len(label)) loss.backward() optimizer.step() pbar.print("%s: %0.2f" % ('train loss', metrics['train_loss'].avg)) pbar.print("%s: %0.2f" % ('train acc', metrics['train_acc'].avg)) # val loop net = net.eval() with torch.no_grad(): for step, item in enumerate(pbar(val_loader)): im = item[0].cuda() label = item[1].cuda() output = net(im) loss = criterion(output, label) accs, _ = accuracy(output, label, topk=(1, )) metrics['val_loss'].update(loss, n=len(label)) metrics['val_acc'].update(accs[0], n=len(label)) net = net.train() # update scheduler scheduler.step(metrics['val_acc'].avg) # send losses to tensorboard for k, v in metrics.items(): pbar.print("Metrics at end of epoch") pbar.print("%s: %0.4f" % (k, v.avg)) writer.add_scalar(k.replace('_', '/'), v.avg, epoch) pbar.print("Learning rate: %0.6f" % optimizer.param_groups[0]['lr']) # do checkpoint as latest util.make_checkpoint(net, optimizer, epoch, metrics['val_acc'].avg, opt.outf, 'latest') if metrics['val_acc'].avg > best_val_acc: pbar.print("Updating best checkpoint at epoch %d" % epoch) pbar.print("Old Best Epoch %d Best Val %0.2f" % (best_val_epoch, best_val_acc)) # do checkpoint as best util.make_checkpoint(net, optimizer, epoch, metrics['val_acc'].avg, opt.outf, 'best') best_val_acc = metrics['val_acc'].avg best_val_epoch = epoch pbar.print("New Best Epoch %d Best Val %0.2f" % (best_val_epoch, best_val_acc)) with open("%s/best_val.txt" % opt.outf, "w") as f: f.write("Best Epoch %d Best Val %0.2f\n" % (best_val_epoch, best_val_acc)) # terminate training if reached min LR and best validation is # not improving if (float(optimizer.param_groups[0]['lr']) <= 1e-6 and epoch >= best_val_epoch + 20): pbar.print("Exiting training") pbar.print("Best Val epoch %d" % best_val_epoch) pbar.print("Curr epoch %d" % epoch) break
def main(opt): # Logging trace_file = os.path.join(opt['output_dir'], '{}_trace.txt'.format(opt['exp_name'])) # Load data if opt['dataset'] == 'cifar-fs': train_data = get_dataset('cifar-fs-train-train', opt['dataroot']) val_data = get_dataset('cifar-fs-val', opt['dataroot']) test_data = get_dataset('cifar-fs-test', opt['dataroot']) tr = get_transform('cifar_resize_normalize') normalize = cifar_normalize elif opt['dataset'] == 'miniimagenet': train_data = get_dataset('miniimagenet-train-train', opt['dataroot']) val_data = get_dataset('miniimagenet-val', opt['dataroot']) test_data = get_dataset('miniimagenet-test', opt['dataroot']) tr = get_transform('cifar_resize_normalize_84') normalize = cifar_normalize if opt['input_regularization'] == 'oe': reg_data = load_ood_data({ 'name': 'tinyimages', 'ood_scale': 1, 'n_anom': 50000, }) if not opt['ooe_only']: if opt['db']: ood_distributions = ['ooe', 'gaussian'] else: ood_distributions = [ 'ooe', 'gaussian', 'rademacher', 'texture3', 'svhn', 'tinyimagenet', 'lsun' ] if opt['input_regularization'] == 'oe': ood_distributions.append('tinyimages') ood_tensors = [('ooe', None)] + [(out_name, load_ood_data({ 'name': out_name, 'ood_scale': 1, 'n_anom': 10000, })) for out_name in ood_distributions[1:]] # Load trained model loaded = torch.load(opt['model.model_path']) if not isinstance(loaded, OrderedDict): fs_model = loaded else: classifier = ResNetClassifier(64, train_data['im_size']).to(device) classifier.load_state_dict(loaded) fs_model = Protonet(classifier.encoder) fs_model.eval() fs_model = fs_model.to(device) # Init Confidence Methods if opt['confidence_method'] == 'oec': init_sample = load_episode(train_data, tr, opt['data.test_way'], opt['data.test_shot'], opt['data.test_query'], device) conf_model = OECConfidence(None, fs_model, init_sample, opt) elif opt['confidence_method'] == 'deep-oec': init_sample = load_episode(train_data, tr, opt['data.test_way'], opt['data.test_shot'], opt['data.test_query'], device) conf_model = DeepOECConfidence(None, fs_model, init_sample, opt) elif opt['confidence_method'] == 'dm-iso': encoder = fs_model.encoder deep_mahala_obj = DeepMahala(None, None, None, encoder, device, num_feats=encoder.depth, num_classes=train_data['n_classes'], pretrained_path="", fit=False, normalize=None) conf_model = DMConfidence(deep_mahala_obj, { 'ls': range(encoder.depth), 'reduction': 'max', 'g_magnitude': .1 }, True, 'iso') if opt['pretrained_oec_path']: conf_model.load_state_dict(torch.load(opt['pretrained_oec_path'])) conf_model.to(device) print(conf_model) optimizer = optim.Adam(conf_model.confidence_parameters(), lr=opt['lr'], weight_decay=opt['wd']) scheduler = StepLR(optimizer, step_size=opt['lrsche_step_size'], gamma=opt['lrsche_gamma']) num_param = sum(p.numel() for p in conf_model.confidence_parameters()) print(f"Learning Confidence, Number of Parameters -- {num_param}") if conf_model.pretrain_parameters() is not None: pretrain_optimizer = optim.Adam(conf_model.pretrain_parameters(), lr=10) pretrain_iter = 100 start_idx = 0 if opt['resume']: last_ckpt_path = os.path.join(opt['output_dir'], 'last_ckpt.pt') if os.path.exists(last_ckpt_path): try: last_ckpt = torch.load(last_ckpt_path) if 'conf_model' in last_ckpt: conf_model = last_ckpt['conf_model'] else: sd = last_ckpt['conf_model_sd'] conf_model.load_state_dict(sd) optimizer = last_ckpt['optimizer'] pretrain_optimizer = last_ckpt['pretrain_optimizer'] scheduler = last_ckpt['scheduler'] start_idx = last_ckpt['outer_idx'] conf_model.to(device) except EOFError: print( "\n\nResuming but got EOF error, starting from init..\n\n") wandb.run.name = opt['exp_name'] wandb.run.save() # try: wandb.watch(conf_model) # except: # resuming a run # pass # Eval and Logging confs = { opt['confidence_method']: conf_model, } if opt['confidence_method'] == 'oec': confs['ed'] = FSCConfidence(fs_model, 'ed') elif opt['confidence_method'] == 'deep-oec': encoder = fs_model.encoder deep_mahala_obj = DeepMahala(None, None, None, encoder, device, num_feats=encoder.depth, num_classes=train_data['n_classes'], pretrained_path="", fit=False, normalize=None) confs['dm'] = DMConfidence(deep_mahala_obj, { 'ls': range(encoder.depth), 'reduction': 'max', 'g_magnitude': 0 }, True, 'iso').to(device) # Temporal Ensemble for Evaluation if opt['n_ensemble'] > 1: nets = [deepcopy(conf_model) for _ in range(opt['n_ensemble'])] confs['mixture-' + opt['confidence_method']] = Ensemble( nets, 'mixture') confs['poe-' + opt['confidence_method']] = Ensemble(nets, 'poe') ensemble_update_interval = opt['eval_every_outer'] // opt['n_ensemble'] iteration_fieldnames = ['global_iteration'] for c in confs: iteration_fieldnames += [ f'{c}_train_ooe', f'{c}_val_ooe', f'{c}_test_ooe', f'{c}_ood' ] iteration_logger = CSVLogger(every=0, fieldnames=iteration_fieldnames, filename=os.path.join(opt['output_dir'], 'iteration_log.csv')) best_val_ooe = 0 PATIENCE = 5 # Number of evaluations to wait waited = 0 progress_bar = tqdm(range(start_idx, opt['train_iter'])) for outer_idx in progress_bar: sample = load_episode(train_data, tr, opt['data.test_way'], opt['data.test_shot'], opt['data.test_query'], device) conf_model.train() if opt['full_supervision']: # sanity check conf_model.support(sample['xs']) in_score = conf_model.score(sample['xq'], detach=False).squeeze() out_score = conf_model.score(sample['ooc_xq'], detach=False).squeeze() out_scores = [out_score] for curr_ood, ood_tensor in ood_tensors: if curr_ood == 'ooe': continue start = outer_idx % (len(ood_tensor) // 2) stop = min( start + sample['xq'].shape[0] * sample['xq'].shape[0], len(ood_tensor) // 2) oxq = torch.stack([tr(x) for x in ood_tensor[start:stop]]).to(device) o = conf_model.score(oxq, detach=False).squeeze() out_scores.append(o) # out_score = torch.cat(out_scores) in_score = in_score.repeat(len(ood_tensors)) loss, acc = compute_loss_bce(in_score, out_score, mean_center=False) else: conf_model.support(sample['xs']) if opt['interpolate']: half_n_way = sample['xq'].shape[0] // 2 interp = .5 * (sample['xq'][:half_n_way] + sample['xq'][half_n_way:2 * half_n_way]) sample['ooc_xq'][:half_n_way] = interp if opt['input_regularization'] == 'oe': # Reshape ooc_xq nw, nq, c, h, w = sample['ooc_xq'].shape sample['ooc_xq'] = sample['ooc_xq'].view(1, nw * nq, c, h, w) oe_bs = int(nw * nq * opt['input_regularization_percent']) start = (outer_idx * oe_bs) % len(reg_data) end = np.min([start + oe_bs, len(reg_data)]) oe_batch = torch.stack([tr(x) for x in reg_data[start:end] ]).to(device) oe_batch = oe_batch.unsqueeze(0) sample['ooc_xq'][:, :oe_batch.shape[1]] = oe_batch if opt['in_out_1_batch']: inps = torch.cat([sample['xq'], sample['ooc_xq']], 1) scores = conf_model.score(inps, detach=False).squeeze() in_score, out_score = scores[:sample['xq'].shape[1]], scores[ sample['xq'].shape[1]:] else: in_score = conf_model.score(sample['xq'], detach=False).squeeze() out_score = conf_model.score(sample['ooc_xq'], detach=False).squeeze() loss, acc = compute_loss_bce(in_score, out_score, mean_center=False) if conf_model.pretrain_parameters( ) is not None and outer_idx < pretrain_iter: pretrain_optimizer.zero_grad() loss.backward() pretrain_optimizer.step() else: optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() progress_bar.set_postfix(loss='{:.3e}'.format(loss), acc='{:.3e}'.format(acc)) # Update Ensemble if opt['n_ensemble'] > 1 and outer_idx % ensemble_update_interval == 0: update_ind = (outer_idx // ensemble_update_interval) % opt['n_ensemble'] if opt['db']: print(f"===> Updating Ensemble: {update_ind}") confs['mixture-' + opt['confidence_method']].nets[update_ind] = deepcopy( conf_model) confs['poe-' + opt['confidence_method']].nets[update_ind] = deepcopy( conf_model) # AUROC eval if outer_idx % opt['eval_every_outer'] == 0: if not opt['eval_in_train']: conf_model.eval() # Eval.. stats_dict = {'global_iteration': outer_idx} for conf_name, conf in confs.items(): conf.eval() # OOE eval ooe_aurocs = {} for split, in_data in [('train', train_data), ('val', val_data), ('test', test_data)]: auroc = np.mean( eval_ood_aurocs( None, in_data, tr, opt['data.test_way'], opt['data.test_shot'], opt['data.test_query'], opt['data.test_episodes'], device, conf, no_grad=False if opt['confidence_method'].startswith('dm') else True)['aurocs']) ooe_aurocs[split] = auroc print_str = '{}, iter: {} ({}), auroc: {:.3e}'.format( conf_name, outer_idx, split, ooe_aurocs[split]) _print_and_log(print_str, trace_file) stats_dict[f'{conf_name}_train_ooe'] = ooe_aurocs['train'] stats_dict[f'{conf_name}_val_ooe'] = ooe_aurocs['val'] stats_dict[f'{conf_name}_test_ooe'] = ooe_aurocs['test'] # OOD eval if not opt['ooe_only']: aurocs = [] for curr_ood, ood_tensor in ood_tensors: auroc = np.mean( eval_ood_aurocs( ood_tensor, test_data, tr, opt['data.test_way'], opt['data.test_shot'], opt['data.test_query'], opt['data.test_episodes'], device, conf, no_grad=False if opt['confidence_method'].startswith('dm') else True)['aurocs']) aurocs.append(auroc) print_str = '{}, iter: {} ({}), auroc: {:.3e}'.format( conf_name, outer_idx, curr_ood, auroc) _print_and_log(print_str, trace_file) mean_ood_auroc = np.mean(aurocs) print_str = '{}, iter: {} (OOD_mean), auroc: {:.3e}'.format( conf_name, outer_idx, mean_ood_auroc) _print_and_log(print_str, trace_file) stats_dict[f'{conf_name}_ood'] = mean_ood_auroc iteration_logger.writerow(stats_dict) plot_csv(iteration_logger.filename, iteration_logger.filename) wandb.log(stats_dict) if stats_dict[f'{opt["confidence_method"]}_val_ooe'] > best_val_ooe: conf_model.cpu() torch.save( conf_model.state_dict(), os.path.join(opt['output_dir'], opt['exp_name'] + '_conf_best.pt')) conf_model.to(device) # Ckpt ensemble if opt['n_ensemble'] > 1: ensemble = confs['mixture-' + opt['confidence_method']] ensemble.cpu() torch.save( ensemble.state_dict(), os.path.join(opt['output_dir'], opt['exp_name'] + '_ensemble_best.pt')) ensemble.to(device) waited = 0 else: waited += 1 if waited >= PATIENCE: print("PATIENCE exceeded...exiting") sys.exit() # For `resume` conf_model.cpu() torch.save( { 'conf_model_sd': conf_model.state_dict(), 'optimizer': optimizer, 'pretrain_optimizer': pretrain_optimizer if conf_model.pretrain_parameters() is not None else None, 'scheduler': scheduler, 'outer_idx': outer_idx, }, os.path.join(opt['output_dir'], 'last_ckpt.pt')) conf_model.to(device) conf_model.train() sys.exit()
# setup output directory if args.stylemix_layer is None: save_dir = 'results/classifiers/cifar10/latentclassifier' else: save_dir = 'results/classifiers/cifar10/latentclassifier_stylemix_%s' % args.stylemix_layer os.makedirs(save_dir, exist_ok=True) pidfile.exit_if_job_done(save_dir) torch.manual_seed(0) 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..') transform_train = data.get_transform('cifar10', 'imtrain') transform_test = data.get_transform('cifar10', 'imval') trainset = data.get_dataset('cifar10', 'train', load_w=True, transform=transform_train) trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2) # using validation partition testset = data.get_dataset('cifar10', 'val', load_w=True, transform=transform_test) testloader = torch.utils.data.DataLoader(testset,
def train(opt): print("Random Seed: ", opt.seed) random.seed(opt.seed) torch.manual_seed(opt.seed) cudnn.benchmark = True device = 'cuda' batch_size = int(opt.batch_size) domain = opt.domain # tensorboard os.makedirs(os.path.join(opt.outf, 'runs'), exist_ok=True) writer = SummaryWriter(log_dir=os.path.join(opt.outf, 'runs')) # datasets train_transform = data.get_transform(domain, 'imtrain') train_dset = data.get_dataset(domain, 'train', load_w=True, transform=train_transform) print("Training transform:") print(train_transform) val_transform = data.get_transform(domain, 'imval') val_dset = data.get_dataset(domain, 'val', load_w=True, transform=val_transform) print("Validation transform:") print(val_transform) train_loader = DataLoader(train_dset, batch_size=opt.batch_size, shuffle=True, pin_memory=False, num_workers=opt.workers) val_loader = DataLoader(val_dset, batch_size=opt.batch_size, shuffle=False, pin_memory=False, num_workers=opt.workers) # classifier: resnet18 model net = torchvision.models.resnet18( num_classes=len(train_dset.coarse_labels)) # load the model weights to finetune from ckpt_path = ('results/classifiers/%s/%s/net_best.pth' % (opt.domain, opt.finetune_from)) print("Finetuning model from %s" % ckpt_path) ckpt = torch.load(ckpt_path) state_dict = ckpt['state_dict'] net.load_state_dict(state_dict) net = net.to(device) # losses + optimizers criterion = nn.CrossEntropyLoss().to(device) # loss(output, target) optimizer = optim.Adam(net.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', patience=10, min_lr=1e-6, verbose=True) start_ep = 0 best_val_acc = 0.0 best_val_epoch = 0 # val tensor transform does not flip, train tensor transform # mimics random resized crop and random flip val_tensor_transform = data.get_transform(domain, 'tensorbase') train_tensor_transform = data.get_transform(domain, 'tensormixedtrain') # load GAN generator = domain_generator.define_generator('stylegan2', domain) for epoch in pbar(range(start_ep, opt.niter + 1)): # average meter for train/val loss and train/val acc metrics = dict(train_loss=util.AverageMeter(), val_loss=util.AverageMeter(), train_acc=util.AverageMeter(), val_acc=util.AverageMeter()) # train loop for step, item in enumerate(pbar(train_loader)): im_orig = item[0].cuda() opt_w = item[1].cuda() label = item[2].cuda() with torch.no_grad(): if opt.perturb_type == 'stylemix': seed = epoch * len(train_loader) + step mix_latent = generator.seed2w(n=opt_w.shape[0], seed=seed) generated_im = generator.perturb_stylemix( opt_w, opt.perturb_layer, mix_latent, n=opt_w.shape[0], is_eval=False) elif opt.perturb_type == 'isotropic': # used minimum value for isotropic setting eps = np.min( generator.perturb_settings['isotropic_eps_%s' % opt.perturb_layer]) generated_im = generator.perturb_isotropic( opt_w, opt.perturb_layer, eps=eps, n=opt_w.shape[0], is_eval=False) elif opt.perturb_type == 'pca': # used median value for pca setting eps = np.median(generator.perturb_settings['pca_eps']) generated_im = generator.perturb_pca(opt_w, opt.perturb_layer, eps=eps, n=opt_w.shape[0], is_eval=False) else: generated_im = generator.decode(opt_w) generated_im = train_tensor_transform(generated_im) # sanity check that the shapes match assert (generated_im.shape == im_orig.shape) # with 50% chance, take the original image for # training the classifier im = im_orig if torch.rand(1) > 0.5 else generated_im net.zero_grad() output = net(im) loss = criterion(output, label) accs, _ = accuracy(output, label, topk=(1, )) metrics['train_loss'].update(loss, n=len(label)) metrics['train_acc'].update(accs[0], n=len(label)) loss.backward() optimizer.step() if step % 200 == 0: pbar.print("%s: %0.6f" % ('train loss', metrics['train_loss'].avg)) pbar.print("%s: %0.6f" % ('train acc', metrics['train_acc'].avg)) # val loop net = net.eval() with torch.no_grad(): for step, item in enumerate(pbar(val_loader)): im_orig = item[0].cuda() opt_w = item[1].cuda() label = item[2].cuda() with torch.no_grad(): # evaluate on the generated image im = generator.decode(opt_w) im = val_tensor_transform(im) assert (im.shape == im_orig.shape) output = net(im) loss = criterion(output, label) accs, _ = accuracy(output, label, topk=(1, )) metrics['val_loss'].update(loss, n=len(label)) metrics['val_acc'].update(accs[0], n=len(label)) net = net.train() scheduler.step(metrics['val_acc'].avg) # send losses to tensorboard for k, v in metrics.items(): pbar.print("Metrics at end of epoch") pbar.print("%s: %0.4f" % (k, v.avg)) writer.add_scalar(k.replace('_', '/'), v.avg, epoch) pbar.print("Learning rate: %0.6f" % optimizer.param_groups[0]['lr']) # do checkpoint as latest util.make_checkpoint(net, optimizer, epoch, metrics['val_acc'].avg, opt.outf, 'latest') if metrics['val_acc'].avg > best_val_acc: pbar.print("Updating best checkpoint at epoch %d" % epoch) pbar.print("Old Best Epoch %d Best Val %0.6f" % (best_val_epoch, best_val_acc)) # do checkpoint as best util.make_checkpoint(net, optimizer, epoch, metrics['val_acc'].avg, opt.outf, 'best') best_val_acc = metrics['val_acc'].avg best_val_epoch = epoch pbar.print("New Best Epoch %d Best Val %0.6f" % (best_val_epoch, best_val_acc)) with open("%s/best_val.txt" % opt.outf, "w") as f: f.write("Best Epoch %d Best Val %0.6f\n" % (best_val_epoch, best_val_acc)) if (float(optimizer.param_groups[0]['lr']) <= 1e-6 and epoch >= best_val_epoch + 50): pbar.print("Exiting training") pbar.print("Best Val epoch %d" % best_val_epoch) pbar.print("Curr epoch %d" % epoch) break
def main(args): device = 'cuda:0' if args['data.cuda'] else 'cpu' args['log.exp_dir'] = args['log.exp_dir'] if not os.path.isdir(args['log.exp_dir']): os.makedirs(args['log.exp_dir']) # save opts with open(os.path.join(args['log.exp_dir'], 'args.json'), 'w') as f: json.dump(args, f) f.write('\n') # Loggin iteration_fieldnames = ['global_iteration', 'val_acc'] iteration_logger = CSVLogger(every=0, fieldnames=iteration_fieldnames, filename=os.path.join(args['log.exp_dir'], 'iteration_log.csv')) # Set the random seed manually for reproducibility. np.random.seed(args['seed']) torch.manual_seed(args['seed']) if args['data.cuda']: torch.cuda.manual_seed(args['seed']) if args['data.dataset'] == 'omniglot': raise train_tr = None test_tr = None elif args['data.dataset'] == 'miniimagenet': train_data = get_dataset('miniimagenet-train-train', args['dataroot']) val_data = get_dataset('miniimagenet-val', args['dataroot']) test_data = get_dataset('miniimagenet-test', args['dataroot']) train_tr = get_transform( 'cifar_augment_normalize_84' if args['data_augmentation'] else 'cifar_normalize') test_tr = get_transform('cifar_normalize') elif args['data.dataset'] == 'cifar100': train_data = get_dataset('cifar-fs-train-train') val_data = get_dataset('cifar-fs-val') test_data = get_dataset('cifar-fs-test') train_tr = get_transform( 'cifar_augment_normalize' if args['data_augmentation'] else 'cifar_normalize') test_tr = get_transform('cifar_normalize') else: raise model = protonet.create_model(**args) # Load model loaded = torch.load(args['model.model_path']) if not 'Protonet' in str(loaded.__class__): pretrained = ResNetClassifier(64, train_data['im_size']).to(device) pretrained.load_state_dict(loaded) model.encoder = pretrained.encoder else: model = loaded model = BaselineFinetune(model.encoder, args['data.way'], args['data.shot'], loss_type='dist') model = model.to(device) def evaluate(data): corrects = [] for _ in tqdm(range(args['data.test_episodes'])): sample = load_episode(data, test_tr, args['data.test_way'], args['data.test_shot'], args['data.test_query'], device) corrects.append(classification_accuracy(sample, model)[0]) acc = torch.mean(torch.cat(corrects)) return acc.item() print("Val acc: ", evaluate(val_data)) print("Test acc: ", evaluate(test_data)) sys.exit()
csv_path = os.path.join(root, model.name) if os.path.isdir(csv_path): os.mkdir(csv_path) csv_path = os.path.join(csv_path, 'tta_' + str(n_crops) + '_sub.csv') else: csv_path = os.path.join(root, 'tta_' + str(n_crops) + '_sub.csv') res.to_csv(csv_path, index=False) return csv_path if __name__ == '__main__': root = './' nums_tta = 7 #7+1 batch_size_test = 64 tsfm_test = get_transform('train') #TTA dataset_test = MyDataset(root=root + '/data', transforms=tsfm_test, phase='test') test_dataloader = torch.utils.data.DataLoader(dataset_test, batch_size=batch_size_test, shuffle=False, num_workers=0) model = Resnet101(pretrained=False) checkpoint_path = os.path.join(root, 'saved_models', 'model_best_' + model.name + '.tar') csv_name = predict(model, root, test_dataloader, batch_size_test,
def main(args): device = 'cuda:0' if args['data.cuda'] else 'cpu' args['log.exp_dir'] = args['log.exp_dir'] if not os.path.isdir(args['log.exp_dir']): os.makedirs(args['log.exp_dir']) # save opts with open(os.path.join(args['log.exp_dir'], 'args.json'), 'w') as f: json.dump(args, f) f.write('\n') # Loggin iteration_fieldnames = ['global_iteration', 'val_acc'] iteration_logger = CSVLogger(every=0, fieldnames=iteration_fieldnames, filename=os.path.join(args['log.exp_dir'], 'iteration_log.csv')) # Set the random seed manually for reproducibility. np.random.seed(args['seed']) torch.manual_seed(args['seed']) if args['data.cuda']: torch.cuda.manual_seed(args['seed']) if args['data.dataset'] == 'omniglot': raise train_tr = None test_tr = None elif args['data.dataset'] == 'miniimagenet': train_data = get_dataset('miniimagenet-train-train', args['dataroot']) val_data = get_dataset('miniimagenet-val', args['dataroot']) train_tr = get_transform( 'cifar_augment_normalize_84' if args['data_augmentation'] else 'cifar_normalize') test_tr = get_transform('cifar_normalize') elif args['data.dataset'] == 'cifar100': train_data = get_dataset('cifar-fs-train-train') val_data = get_dataset('cifar-fs-val') train_tr = get_transform( 'cifar_augment_normalize' if args['data_augmentation'] else 'cifar_normalize') test_tr = get_transform('cifar_normalize') else: raise model = protonet.create_model(**args) if args['model.model_path'] != '': loaded = torch.load(args['model.model_path']) if not 'Protonet' in str(loaded.__class__): pretrained = ResNetClassifier(64, train_data['im_size']).to(device) pretrained.load_state_dict(loaded) model.encoder = pretrained.encoder else: model = loaded model = model.to(device) max_epoch = args['train.epochs'] epoch = 0 stop = False patience_elapsed = 0 best_metric_value = 0.0 def evaluate(): nonlocal best_metric_value nonlocal patience_elapsed nonlocal stop nonlocal epoch corrects = [] for _ in tqdm(range(args['data.test_episodes']), desc="Epoch {:d} Val".format(epoch + 1)): sample = load_episode(val_data, test_tr, args['data.test_way'], args['data.test_shot'], args['data.test_query'], device) corrects.append(classification_accuracy(sample, model)[0]) val_acc = torch.mean(torch.cat(corrects)) iteration_logger.writerow({ 'global_iteration': epoch, 'val_acc': val_acc.item() }) plot_csv(iteration_logger.filename, iteration_logger.filename) print(f"Epoch {epoch}: Val Acc: {val_acc}") if val_acc > best_metric_value: best_metric_value = val_acc print("==> best model (metric = {:0.6f}), saving model...".format( best_metric_value)) model.cpu() torch.save(model, os.path.join(args['log.exp_dir'], 'best_model.pt')) model.to(device) patience_elapsed = 0 else: patience_elapsed += 1 if patience_elapsed > args['train.patience']: print("==> patience {:d} exceeded".format( args['train.patience'])) stop = True optim_method = getattr(optim, args['train.optim_method']) params = model.parameters() optimizer = optim_method(params, lr=args['train.learning_rate'], weight_decay=args['train.weight_decay']) scheduler = lr_scheduler.StepLR(optimizer, args['train.decay_every'], gamma=0.5) while epoch < max_epoch and not stop: evaluate() model.train() if epoch % args['ckpt_every'] == 0: model.cpu() torch.save(model, os.path.join(args['log.exp_dir'], f'model_{epoch}.pt')) model.to(device) scheduler.step() for _ in tqdm(range(args['data.train_episodes']), desc="Epoch {:d} train".format(epoch + 1)): sample = load_episode(train_data, train_tr, args['data.way'], args['data.shot'], args['data.query'], device) optimizer.zero_grad() loss, output = model.loss(sample) loss.backward() optimizer.step() epoch += 1
torch.manual_seed(args.seed) np.random.seed(args.seed) random.seed(args.seed) # lock the experiment directory lockdir = f'results/evaluations/{args.domain}/lockfiles/{args.classifier_name}_{args.partition}/image_ensemble_{args.aug_type}' os.makedirs(lockdir, exist_ok=True) pidfile.exit_if_job_done(lockdir, redo=False) # data output filename data_filename = lockdir.replace('lockfiles', 'output') + '.npz' os.makedirs(os.path.dirname(data_filename), exist_ok=True) print("saving result in: %s" % data_filename) # load dataset and classifier val_transform = data.get_transform(args.domain, 'imval') ensemble_transform = data.get_transform(args.domain, args.aug_type) transform = ImageEnsemble(val_transform, ensemble_transform, args.n_ens) print("Ensemble transform:") print(ensemble_transform) if 'celebahq' in args.domain: # for celebahq, load the attribute-specific dataset attribute = args.classifier_name.split('__')[0] dset = data.get_dataset(args.domain, args.partition, attribute, load_w=False, transform=transform) else: dset = data.get_dataset(args.domain, args.partition,
parser.add_argument('--output', type=str, default="output", help="output folder") args = parser.parse_args() # Create directory to store weights if not os.path.exists(args.output): os.makedirs(args.output) model = get_model(args.checkpoint) device = model_device() model.eval() totensor = get_transform(train=False) toimage = transforms.ToPILImage() image_filenames = sorted(glob.glob(args.input)) progress_bar = tqdm(total=len(image_filenames)) for index, filename in enumerate(image_filenames): progress_bar.update(1) image = Image.open(filename).convert("RGB") input_image = image.resize((320, 320)) input_tensor = totensor(input_image).unsqueeze(0).to(device) with torch.no_grad(): output_tensor = model(input_tensor)
def main(opt): eval_exp_name = opt['exp_name'] device = 'cuda:0' # Load data if opt['dataset'] == 'cifar-fs': train_data = get_dataset('cifar-fs-train-train', opt['dataroot']) val_data = get_dataset('cifar-fs-val', opt['dataroot']) test_data = get_dataset('cifar-fs-test', opt['dataroot']) tr = get_transform('cifar_resize_normalize') normalize = cifar_normalize elif opt['dataset'] == 'miniimagenet': train_data = get_dataset('miniimagenet-train-train', opt['dataroot']) val_data = get_dataset('miniimagenet-val', opt['dataroot']) test_data = get_dataset('miniimagenet-test', opt['dataroot']) tr = get_transform('cifar_resize_normalize_84') normalize = cifar_normalize np.random.seed(1234) torch.manual_seed(1234) torch.cuda.manual_seed(1234) if opt['db']: ood_distributions = ['ooe', 'gaussian'] else: ood_distributions = ['ooe', 'gaussian', 'svhn'] # ood_distributions = ['ooe', 'gaussian', 'rademacher', 'texture3', 'svhn','tinyimagenet','lsun'] ood_tensors = [('ooe', None)] + [(out_name, load_ood_data({ 'name': out_name, 'ood_scale': 1, 'n_anom': 10000, })) for out_name in ood_distributions[1:]] # Load trained model loaded = torch.load(opt['model.model_path']) if not isinstance(loaded, OrderedDict): protonet = loaded else: classifier = ResNetClassifier(64, train_data['im_size']).to(device) classifier.load_state_dict(loaded) protonet = Protonet(classifier.encoder) encoder = protonet.encoder encoder.eval() encoder.to(device) protonet.eval() protonet.to(device) # Init Confidence model if opt['ood_method'] == 'deep-ed-iso': deep_mahala_obj = DeepMahala(None, None, None, encoder, device, num_feats=encoder.depth, num_classes=train_data['n_classes'], pretrained_path="", fit=False, normalize=None) conf = DMConfidence(deep_mahala_obj, { 'ls': range(encoder.depth), 'reduction': 'max', 'g_magnitude': 0 }, True, 'iso').to(device) elif opt['ood_method'] == 'native-spp': conf = FSCConfidence(protonet, 'spp') elif opt['ood_method'] == 'oec': oec_opt = json.load( open(os.path.join(os.path.dirname(opt['oec_path']), 'args.json'), 'r')) init_sample = load_episode(train_data, tr, oec_opt['data.test_way'], oec_opt['data.test_shot'], oec_opt['data.test_query'], device) if oec_opt['confidence_method'] == 'oec': oec_conf = OECConfidence(None, protonet, init_sample, oec_opt) else: oec_conf = DeepOECConfidence(None, protonet, init_sample, oec_opt) oec_conf.load_state_dict(torch.load(opt['oec_path'])) oec_conf.eval() oec_conf.to(device) conf = oec_conf # Turn confidence score into a threshold based classifier # Select threshold by "max-accuracy" # Select temperature by "best-calibration" in the binary problem # done using the meta-train set in_scores = [] out_scores = [] for n in tqdm(range(100)): sample = load_episode(train_data, tr, opt['data.test_way'], opt['data.test_shot'], opt['data.test_query'], device) in_score, out_score = score_batch(conf, sample) in_scores.append(in_score) out_scores.append(out_score) in_scores = torch.cat(in_scores) out_scores = torch.cat(out_scores) def _compute_acc(in_scores, out_scores, t): N = len(in_scores) + len(out_scores) return (torch.sum(in_scores >= t) + torch.sum(out_scores < t)).item() / float(N) best_threshold = torch.min(in_scores) best_acc = _compute_acc(in_scores, out_scores, best_threshold) for t in in_scores: acc = _compute_acc(in_scores, out_scores, t) if acc > best_acc: best_acc = acc best_threshold = t def _compute_confs(in_scores, out_scores, t, temp): in_p = torch.sigmoid((in_scores - t) / temp) corrects = in_p >= .5 confs = torch.max(torch.stack([in_p, 1 - in_p]), 0)[0] out_p = torch.sigmoid((out_scores - t) / temp) corrects = torch.cat([corrects, out_p < .5]) confs = torch.cat( [confs, torch.max(torch.stack([out_p, 1 - out_p]), 0)[0]]) return confs, corrects def compute_eces(candidate_temps, in_scores, out_scores, best_threshold): eces = [] for temp in candidate_temps: confs, corrects = _compute_confs(in_scores, out_scores, best_threshold, temp) ece = compute_ece( *prep_accs(confs.numpy(), corrects.numpy(), bins=20)) eces.append(ece) return eces min_log_temp = -1 log_interval = 2 npts = 10 for _ in range(opt['max_temp_select_iter']): print("..selecting temperature") candidate_temps = np.logspace(min_log_temp, min_log_temp + log_interval, npts) eces = compute_eces(candidate_temps, in_scores, out_scores, best_threshold) min_idx = np.argmin(eces) if min_idx == 0: min_log_temp -= log_interval // 2 elif min_idx == npts - 1: min_log_temp += log_interval // 2 else: break best_ece = eces[min_idx] best_temp = candidate_temps[min_idx] print( f"Best ACC:{best_acc}, thresh:{best_threshold}, Best ECE:{best_ece}, temp:{best_temp}" ) def get_95_percent_ci(std): """Computes the 95% confidence interval from the standard deviation.""" return std * 1.96 / np.sqrt(data_opt['data.test_episodes']) active_supervised = defaultdict(list) active_augmented = defaultdict(list) ssl_soft = defaultdict(list) ssl_hard = defaultdict(list) # for ood_idx, curr_ood in tqdm(enumerate(all_distributions)): for curr_ood, ood_tensor in ood_tensors: in_scores = defaultdict(list) out_scores = defaultdict(list) # Compute and collect scores for all examples aurocs, auprs, fprs = defaultdict(list), defaultdict( list), defaultdict(list) for n in tqdm(range(opt['data.test_episodes'])): n_total_query = np.max([ opt['data.test_query'] + opt['n_unlabeled_per_class'], opt['n_distractor_per_class'] ]) sample = load_episode(test_data, tr, opt['data.test_way'], opt['data.test_shot'], n_total_query, device) if curr_ood != 'ooe': bs = opt['data.test_way'] * opt['data.test_query'] ridx = np.random.permutation(bs) sample['ooc_xq'] = torch.stack( [tr(x) for x in ood_tensor[ridx]]).to(device) way, _, c, h, w = sample['xq'].shape sample['ooc_xq'] = sample['ooc_xq'].reshape(way, -1, c, h, w) # if curr_ood in ['gaussian', 'rademacher']: # sample['ooc_xq'] *= 4 all_xq = sample['xq'].clone() sample['xq'] = all_xq[:, :opt[ 'n_unlabeled_per_class']] # Unlabelled pool sample[ 'xq2'] = all_xq[:, opt['n_unlabeled_per_class']: opt['n_unlabeled_per_class'] + opt['data.test_query']] # Final test queries sample['ooc_xq'] = sample[ 'ooc_xq'][:, :opt['n_distractor_per_class']] """ 1. OOD classification on the 'unlabelled' set """ # In vs Out in_score, out_score = score_batch(conf, sample) num_in = in_score.shape[0] confs, corrects = _compute_confs(in_score, out_score, best_threshold, best_temp) in_mask = corrects[:num_in].reshape( sample['xq'].size(0), sample['xq'].size(1)).float().to(device) out_mask = 1 - corrects[num_in:].reshape( sample['ooc_xq'].size(0), sample['ooc_xq'].size(1)).float().to(device) """ 2.0 """ budget_active = in_score.size(0) scores = torch.cat([in_score, out_score], -1) ipdb.set_trace() selected_inds = torch.sort(scores)[1][scores.size(0) - budget_active:] selected_inds_in = selected_inds[selected_inds < in_score.size(0)] budget_mask = torch.zeros(in_score.size(0)).to(device) budget_mask.scatter_(0, selected_inds_in.to(device).long(), 1) budget_mask = budget_mask.reshape( sample['xq'].size(0), sample['xq'].size(1)).float().to(device) """ 2. Add labels to the predicted unlabelled examples """ # Collect the incorrectly kept OOD examples included_distractors = sample['ooc_xq'][out_mask.byte()] # Pad them to N-way multiples, and assign random labels (done simply by reshaping) n_way = sample['xs'].shape[0] im_shape = list(sample['xs'].shape[2:]) n_res = n_way - (included_distractors.shape[0] % n_way) distractor_mask = torch.ones([included_distractors.shape[0] ]).to(device) zeros = torch.zeros([n_res] + im_shape).to(device) included_distractors = torch.cat([included_distractors, zeros]) distractor_mask = torch.cat( [distractor_mask, torch.zeros([n_res]).to(device)]) # the reason we permute is to spread the padded zero across ways included_distractors = included_distractors.reshape( [-1, n_way] + im_shape).permute(1, 0, 2, 3, 4) distractor_mask = distractor_mask.reshape([-1, n_way]).permute(1, 0) """ 2.5 SSL """ # predict k-way using classifier n_way, n_aug_shot, n_ch, n_dim, _ = sample['xq'].shape lpy_dic = protonet.log_p_y(sample['xs'], sample['xq'], mask=None) log_p_y, target_inds = lpy_dic['log_p_y'], lpy_dic['target_inds'] preds = log_p_y.max(-1)[1] def reorder(unlabelled, preds, py, make_soft=True): if py is not None: reshaped_py = py.reshape(-1) n_way, n_aug_shot, n_ch, n_dim, _ = unlabelled.shape reshaped_unlabelled = unlabelled.reshape( n_aug_shot * n_way, n_ch, n_dim, n_dim) reshaped_predicted_labels = preds.reshape(-1) unlabelled = torch.zeros( (n_way, n_aug_shot * n_way, n_ch, n_dim, n_dim)) mask = torch.zeros((n_way, n_aug_shot * n_way)) for idx, label in enumerate(reshaped_predicted_labels): unlabelled[label, idx] = reshaped_unlabelled[idx] # (n_shot, ...) if make_soft: mask[label, idx] = reshaped_py[idx] else: mask[label, idx] = 1 # (n_shot, ) return unlabelled.to(device), mask.to(device) gt_in_unlabelled, gt_in_weights = reorder(sample['xq'], preds, log_p_y.max(-1)[0].exp(), True) _, in_mask_reordered = reorder(sample['xq'], preds, in_mask, True) # for the gt OOD ones lpy_dic = protonet.log_p_y(sample['xs'], sample['ooc_xq'], mask=None) log_p_y = lpy_dic['log_p_y'] preds = log_p_y.max(-1)[1] gt_ood_unlabelled, gt_ood_weights = reorder( sample['ooc_xq'], preds, log_p_y.max(-1)[0].exp(), True) _, out_mask_reordered = reorder(sample['ooc_xq'], preds, out_mask, True) # Support + ALL unlabelled _ssl_soft = compute_acc( protonet, torch.cat([sample['xs'], gt_in_unlabelled, gt_ood_unlabelled], 1), sample['xq2'], torch.cat([ torch.ones(sample['xs'].shape[:2]).to(device), gt_in_weights, gt_ood_weights ], 1)) _acc_hard = compute_acc( protonet, torch.cat([sample['xs'], gt_in_unlabelled, gt_ood_unlabelled], 1), sample['xq2'], torch.cat([ torch.ones(sample['xs'].shape[:2]).to(device), in_mask_reordered * gt_in_weights, out_mask_reordered * gt_ood_weights ], 1)) """ 3. Evaluate k-way accuracy after adding examples """ _active_supervised = compute_acc(protonet, sample['xs'], sample['xq2'], None) # Support + Budgeted unlabelled _active_augmented = compute_acc( protonet, torch.cat([sample['xs'], sample['xq']], 1), sample['xq2'], torch.cat([ torch.ones(sample['xs'].shape[:2]).to(device), budget_mask ], 1)) ssl_soft[curr_ood].append(_ssl_soft) ssl_hard[curr_ood].append(_acc_hard) active_supervised[curr_ood].append(_active_supervised) active_augmented[curr_ood].append(_active_augmented) if not os.path.exists(opt['output_dir']): os.makedirs(opt['output_dir']) pickle.dump((ssl_soft, ssl_hard, active_supervised, active_augmented), open( os.path.join(opt['output_dir'], f'eval_active_{eval_exp_name}.pkl'), 'wb')) print("===> Aggregating results") aggr_args = namedtuple('Arg', ('exp_dir', 'f_acq'))(exp_dir=opt['output_dir'], f_acq='conv4') aggregate_eval_active.main(aggr_args) print('===> Done') sys.exit()
def train(opt): print("Random Seed: ", opt.seed) random.seed(opt.seed) torch.manual_seed(opt.seed) cudnn.benchmark = True device = 'cuda' batch_size = int(opt.batch_size) # tensorboard os.makedirs(os.path.join(opt.outf, 'runs'), exist_ok=True) writer = SummaryWriter(log_dir=os.path.join(opt.outf, 'runs')) # classifier follows architecture and initialization from attribute # classifiers in stylegan: # https://github.com/NVlabs/stylegan/blob/master/metrics/linear_separability.py#L136 # https://github.com/NVlabs/stylegan/blob/master/training/networks_stylegan.py#L564 net = attribute_classifier.D(3, resolution=256, fixed_size=True, use_mbstd=False) # use random normal bc wscale will rescale weights, see tf source # https://github.com/NVlabs/stylegan/blob/master/training/networks_stylegan.py#L148 netinit.init_weights(net, init_type='normal', gain=1.) net = net.to(device) # losses + optimizers bce_loss = nn.BCEWithLogitsLoss() optimizer = optim.Adam(net.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) # datasets -- added random horizonal flipping for training train_transform = data.get_transform('celebahq', 'imtrain') train_dset = data.get_dataset('celebahq', 'train', opt.attribute, load_w=False, transform=train_transform) print("Training transform:") print(train_transform) val_transform = data.get_transform('celebahq', 'imval') val_dset = data.get_dataset('celebahq', 'val', opt.attribute, load_w=False, transform=val_transform) print("Validation transform:") print(val_transform) train_loader = DataLoader(train_dset, batch_size=opt.batch_size, shuffle=True, pin_memory=False, num_workers=opt.workers) val_loader = DataLoader(val_dset, batch_size=opt.batch_size, shuffle=False, pin_memory=False, num_workers=opt.workers) start_ep = 0 best_val_acc = 0.0 best_val_epoch = 0 for epoch in pbar(range(start_ep, opt.niter + 1)): # average meter for train/val loss and train/val acc metrics = dict(train_loss=util.AverageMeter(), val_loss=util.AverageMeter(), train_acc=util.AverageMeter(), val_acc=util.AverageMeter()) # train loop for step, (im, label) in enumerate(pbar(train_loader)): im = im.cuda() label = label.cuda().float() net.zero_grad() logit, softmaxed = attribute_utils.get_softmaxed(net, im) # enforces that negative logit --> our label = 1 loss = bce_loss(logit, 1 - label) predicted = (softmaxed > 0.5).long() correct = (predicted == label).float().mean().item() metrics['train_loss'].update(loss, n=len(label)) metrics['train_acc'].update(correct, n=len(label)) loss.backward() optimizer.step() if step % 200 == 0: pbar.print("%s: %0.2f" % ('train loss', metrics['train_loss'].avg)) pbar.print("%s: %0.2f" % ('train acc', metrics['train_acc'].avg)) # val loop net = net.eval() with torch.no_grad(): for step, (im, label) in enumerate(pbar(val_loader)): im = im.cuda() label = label.cuda().float() logit, softmaxed = attribute_utils.get_softmaxed(net, im) predicted = (softmaxed > 0.5).long() correct = (predicted == label).float().mean().item() loss = bce_loss(logit, 1 - label) metrics['val_loss'].update(loss, n=len(label)) metrics['val_acc'].update(correct, n=len(label)) net = net.train() # send losses to tensorboard for k, v in metrics.items(): pbar.print("Metrics at end of epoch") pbar.print("%s: %0.4f" % (k, v.avg)) writer.add_scalar(k.replace('_', '/'), v.avg, epoch) # do checkpoint as latest util.make_checkpoint(net, optimizer, epoch, metrics['val_acc'].avg, opt.outf, 'latest') if metrics['val_acc'].avg > best_val_acc: pbar.print("Updating best checkpoint at epoch %d" % epoch) pbar.print("Old Best Epoch %d Best Val %0.2f" % (best_val_epoch, best_val_acc)) # do checkpoint as best util.make_checkpoint(net, optimizer, epoch, metrics['val_acc'].avg, opt.outf, 'best') best_val_acc = metrics['val_acc'].avg best_val_epoch = epoch pbar.print("New Best Epoch %d Best Val %0.2f" % (best_val_epoch, best_val_acc)) with open("%s/best_val.txt" % opt.outf, "w") as f: f.write("Best Epoch %d Best Val %0.2f\n" % (best_val_epoch, best_val_acc)) if epoch >= best_val_epoch + 5: pbar.print("Exiting training") pbar.print("Best Val epoch %d" % best_val_epoch) pbar.print("Curr epoch %d" % epoch) break
exit(0) # lock the experiment directory lockdir = f'results/evaluations/{args.domain}/lockfiles/{args.classifier_name}_{args.partition}/gan_ensemble_{args.aug_type}_{args.aug_layer}' if args.apply_tensor_transform: lockdir += '_tensortransform' os.makedirs(lockdir, exist_ok=True) pidfile.exit_if_job_done(lockdir, redo=False) # data output filename data_filename = lockdir.replace('lockfiles', 'output') + '.npz' os.makedirs(os.path.dirname(data_filename), exist_ok=True) print("saving result in: %s" % data_filename) # load dataset and classifier val_transform = data.get_transform(args.domain, 'imval') if 'celebahq' in args.domain: # for celebahq, load the attribute-specific dataset attribute = args.classifier_name.split('__')[0] dset = data.get_dataset(args.domain, args.partition, attribute, load_w=True, transform=val_transform) else: dset = data.get_dataset(args.domain, args.partition, load_w=True, transform=val_transform) loader = DataLoader(dset, batch_size=1,