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,
