def get_samples(target, nb_class=10, sample_index=0):
'''
params:
target : [mnist, cifar10]
nb_class : number of classes
example_index : index of image by class
returns:
original_images (numpy array) : Original images, shape = (number of class, W, H, C)
pre_images (torch array) : Preprocessing images, shape = (number of class, C, W, H)
target_classes (dictionary) : keys = class index, values = class name
model (pytorch model) : pretrained model
'''
if target == 'mnist':
image_size = (28, 28, 1)
_, _, testloader = mnist_load()
testset = testloader.dataset
elif target == 'cifar10':
image_size = (32, 32, 3)
_, _, testloader = cifar10_load()
testset = testloader.dataset
# idx2class
target_class2idx = testset.class_to_idx
target_classes = dict(
zip(list(target_class2idx.values()), list(target_class2idx.keys())))
# select images
idx_by_class = [
np.where(np.array(testset.targets) == i)[0][sample_index]
for i in range(nb_class)
]
original_images = testset.data[idx_by_class]
if not isinstance(original_images, np.ndarray):
original_images = original_images.numpy()
original_images = original_images.reshape((nb_class, ) + image_size)
# select targets
if isinstance(testset.targets, list):
original_targets = torch.LongTensor(testset.targets)[idx_by_class]
else:
original_targets = testset.targets[idx_by_class]
# model load
weights = torch.load('../checkpoint/simple_cnn_{}.pth'.format(target))
model = SimpleCNN(target)
model.load_state_dict(weights['model'])
# image preprocessing
pre_images = torch.zeros(original_images.shape)
pre_images = np.transpose(pre_images, (0, 3, 1, 2))
for i in range(len(original_images)):
pre_images[i] = testset.transform(original_images[i])
return original_images, original_targets, pre_images, target_classes, model
# TODO: Tensorboard Check
# python main.py --train --target=['mnist','cifar10'] --attention=['CAM','CBAM','RAN','WARN']
if args.train:
main(args=args)
elif args.eval == 'selectivity':
# make evalutation directory
if not os.path.isdir('../evaluation'):
os.mkdir('../evaluation')
# pretrained model load
weights = torch.load('../checkpoint/simple_cnn_{}.pth'.format(
args.target))
model = SimpleCNN(args.target)
model.load_state_dict(weights['model'])
# selectivity evaluation
selectivity_method = Selectivity(model=model,
target=args.target,
batch_size=args.batch_size,
method=args.method,
sample_pct=args.ratio)
# evaluation
selectivity_method.eval(steps=args.steps, savedir='../evaluation')
elif (args.eval == 'ROAR') or (args.eval == 'KAR'):
# ratio
ratio_lst = np.arange(0, 1, args.ratio)[1:] # exclude zero
for ratio in ratio_lst:
main(args=args, ratio=ratio)
