67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
import os
|
|
os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
|
|
import argparse
|
|
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
import pandas as pd
|
|
from scipy import stats
|
|
from src.utils.utilities import *
|
|
from src.metrics.swap import SWAP
|
|
from src.datasets.utilities import get_datasets
|
|
from src.search_space.networks import *
|
|
|
|
# Settings for console outputs
|
|
import warnings
|
|
warnings.simplefilter(action='ignore', category=FutureWarning)
|
|
warnings.simplefilter(action='ignore', category=UserWarning)
|
|
|
|
parser = argparse.ArgumentParser()
|
|
|
|
# general setting
|
|
parser.add_argument('--data_path', default="datasets", type=str, nargs='?', help='path to the image dataset (datasets or datasets/ILSVRC/Data/CLS-LOC)')
|
|
parser.add_argument('--seed', default=0, type=int, help='random seed')
|
|
parser.add_argument('--device', default="mps", type=str, nargs='?', help='setup device (cpu, mps or cuda)')
|
|
parser.add_argument('--repeats', default=32, type=int, nargs='?', help='times of calculating the training-free metric')
|
|
parser.add_argument('--input_samples', default=16, type=int, nargs='?', help='input batch size for training-free metric')
|
|
|
|
args = parser.parse_args()
|
|
|
|
if __name__ == "__main__":
|
|
|
|
device = torch.device(args.device)
|
|
|
|
arch_info = pd.read_csv(args.data_path+'/DARTS_archs_CIFAR10.csv', names=['genotype', 'valid_acc'], sep=',')
|
|
|
|
train_data, _, _ = get_datasets('cifar10', args.data_path, (args.input_samples, 3, 32, 32), -1)
|
|
train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.input_samples, num_workers=0, pin_memory=True)
|
|
loader = iter(train_loader)
|
|
inputs, _ = next(loader)
|
|
|
|
results = []
|
|
|
|
for index, i in arch_info.iterrows():
|
|
print(f'Evaluating network: {index}')
|
|
|
|
network = Network(3, 10, 1, eval(i.genotype))
|
|
network = network.to(device)
|
|
|
|
swap = SWAP(model=network, inputs=inputs, device=device, seed=args.seed)
|
|
|
|
swap_score = []
|
|
|
|
for _ in range(args.repeats):
|
|
network = network.apply(network_weight_gaussian_init)
|
|
swap.reinit()
|
|
swap_score.append(swap.forward())
|
|
swap.clear()
|
|
|
|
results.append([np.mean(swap_score), i.valid_acc])
|
|
|
|
results = pd.DataFrame(results, columns=['swap_score', 'valid_acc'])
|
|
print()
|
|
print(f'Spearman\'s Correlation Coefficient: {stats.spearmanr(results.swap_score, results.valid_acc)[0]}')
|
|
|
|
|
|
|