145 lines
4.9 KiB
Python
145 lines
4.9 KiB
Python
import os
|
|
import argparse
|
|
import random
|
|
import numpy as np
|
|
|
|
import matplotlib.pyplot as plt
|
|
from datasets import get_datasets
|
|
from config_utils import load_config
|
|
|
|
from nas_201_api import NASBench201API as API
|
|
from models import get_cell_based_tiny_net
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
def get_batch_jacobian(net, data_loader, device):
|
|
data_iterator = iter(data_loader)
|
|
x, target = next(data_iterator)
|
|
x = x.to(device)
|
|
net.zero_grad()
|
|
x.requires_grad_(True)
|
|
_, y = net(x)
|
|
y.backward(torch.ones_like(y))
|
|
jacob = x.grad.detach()
|
|
return jacob, target.detach()
|
|
|
|
def plot_hist(jacob, ax, colour):
|
|
xx = jacob.reshape(jacob.size(0), -1).cpu().numpy()
|
|
corrs = np.corrcoef(xx)
|
|
ax.hist(corrs.flatten(), bins=100, color=colour)
|
|
|
|
def decide_plot(acc, plt_cts, num_rows, boundaries=[60., 70., 80., 90.]):
|
|
if acc < boundaries[0]:
|
|
plt_col = 0
|
|
accrange = f'< {boundaries[0]}%'
|
|
elif acc < boundaries[1]:
|
|
plt_col = 1
|
|
accrange = f'[{boundaries[0]}% , {boundaries[1]}%)'
|
|
elif acc < boundaries[2]:
|
|
plt_col = 2
|
|
accrange = f'[{boundaries[1]}% , {boundaries[2]}%)'
|
|
elif acc < boundaries[3]:
|
|
accrange = f'[{boundaries[2]}% , {boundaries[3]}%)'
|
|
plt_col = 3
|
|
else:
|
|
accrange = f'>= {boundaries[3]}%'
|
|
plt_col = 4
|
|
|
|
can_plot = False
|
|
plt_row = 0
|
|
if plt_cts[plt_col] < num_rows:
|
|
can_plot = True
|
|
plt_row = plt_cts[plt_col]
|
|
plt_cts[plt_col] += 1
|
|
|
|
return can_plot, plt_row, plt_col, accrange
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description='Plot histograms of correlation matrix')
|
|
parser.add_argument('--data_loc', default='../datasets/cifar/', type=str, help='dataset folder')
|
|
parser.add_argument('--api_loc', default='../datasets/NAS-Bench-201-v1_1-096897.pth',
|
|
type=str, help='path to API')
|
|
parser.add_argument('--arch_start', default=0, type=int)
|
|
parser.add_argument('--arch_end', default=15625, type=int)
|
|
parser.add_argument('--seed', default=42, type=int)
|
|
parser.add_argument('--GPU', default='0', type=str)
|
|
parser.add_argument('--batch_size', default=256, type=int)
|
|
|
|
args = parser.parse_args()
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = args.GPU
|
|
|
|
# Reproducibility
|
|
torch.backends.cudnn.deterministic = True
|
|
torch.backends.cudnn.benchmark = False
|
|
random.seed(args.seed)
|
|
np.random.seed(args.seed)
|
|
torch.manual_seed(args.seed)
|
|
|
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|
|
|
ARCH_START = args.arch_start
|
|
ARCH_END = args.arch_end
|
|
|
|
criterion = nn.CrossEntropyLoss()
|
|
train_data, valid_data, xshape, class_num = get_datasets('cifar10', args.data_loc, 0)
|
|
|
|
cifar_split = load_config('config_utils/cifar-split.txt', None, None)
|
|
train_split, valid_split = cifar_split.train, cifar_split.valid
|
|
train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size,
|
|
num_workers=0, pin_memory=True, sampler= torch.utils.data.sampler.SubsetRandomSampler(train_split))
|
|
|
|
scores = []
|
|
accs = []
|
|
|
|
plot_shape = (25, 5)
|
|
num_plots = plot_shape[0]*plot_shape[1]
|
|
fig, axes = plt.subplots(*plot_shape, sharex=True, figsize=(9, 9) )
|
|
plt_cts = [0 for i in range(plot_shape[1])]
|
|
|
|
api = API(args.api_loc)
|
|
|
|
archs = list(range(ARCH_START, ARCH_END))
|
|
colours = ['#811F41', '#A92941', '#D15141', '#EF7941', '#F99C4B']
|
|
|
|
strs = []
|
|
random.shuffle(archs)
|
|
for arch in archs:
|
|
try:
|
|
config = api.get_net_config(arch, 'cifar10')
|
|
archinfo = api.query_meta_info_by_index(arch)
|
|
acc = archinfo.get_metrics('cifar10-valid', 'x-valid')['accuracy']
|
|
|
|
network = get_cell_based_tiny_net(config)
|
|
network = network.to(device)
|
|
jacobs, labels = get_batch_jacobian(network, train_loader, device)
|
|
|
|
boundaries = [60., 70., 80., 90.]
|
|
can_plt, row, col, accrange = decide_plot(acc, plt_cts, plot_shape[0], boundaries)
|
|
if not can_plt:
|
|
continue
|
|
axes[row, col].axis('off')
|
|
|
|
plot_hist(jacobs, axes[row, col], colours[col])
|
|
if row == 0:
|
|
axes[row, col].set_title(f'{accrange}')
|
|
|
|
if row + 1 == plot_shape[0]:
|
|
axes[row, col].axis('on')
|
|
plt.setp(axes[row, col].get_xticklabels(), fontsize=12)
|
|
axes[row, col].spines["top"].set_visible(False)
|
|
axes[row, col].spines["right"].set_visible(False)
|
|
axes[row, col].spines["left"].set_visible(False)
|
|
axes[row, col].set_yticks([])
|
|
|
|
if sum(plt_cts) == num_plots:
|
|
plt.tight_layout()
|
|
plt.savefig(f'results/histograms_cifar10val_batch{args.batch_size}.png')
|
|
plt.show()
|
|
break
|
|
except Exception as e:
|
|
plt_cts[col] -= 1
|
|
continue
|