Spaces:
Sleeping
Sleeping
| from src.search import GeneticSearch | |
| from src.hw_nats_fast_interface import HW_NATS_FastInterface | |
| from src.utils import DEVICES, union_of_dicts | |
| import numpy as np | |
| import argparse | |
| import json | |
| def parse_args()->object: | |
| """Args function. | |
| Returns: | |
| (object): args parser | |
| """ | |
| parser = argparse.ArgumentParser() | |
| # this selects the dataset to be considered for the search | |
| parser.add_argument( | |
| "--dataset", | |
| default="cifar10", | |
| type=str, | |
| help="Dataset to be considered. One in ['cifar10', 'cifar100', 'ImageNet16-120'].s", | |
| choices=["cifar10", "cifar100", "ImageNet16-120"] | |
| ) | |
| # this selects the target device to be considered for the search | |
| parser.add_argument( | |
| "--device", | |
| default="edgegpu", | |
| type=str, | |
| help="Device to be considered. One in ['edgegpu', 'eyeriss', 'fpga'].", | |
| choices=["edgegpu", "eyeriss", "fpga"] | |
| ) | |
| # when this flag is triggered, the search is hardware-agnostic (penalized with FLOPS and params) | |
| parser.add_argument("--device-agnostic", action="store_true", help="Flag to trigger hardware-agnostic search.") | |
| parser.add_argument("--n-generations", default=50, type=int, help="Number of generations to let the genetic algorithm run.") | |
| parser.add_argument("--n-runs", default=30, type=int, help="Number of runs used to compute the average test accuracy.") | |
| parser.add_argument("--performance-weight", default=0.65, type=float, help="Weight of the performance metric in the fitness function.") | |
| parser.add_argument("--hardware-weight", default=0.35, type=float, help="Weight of the hardware metric in the fitness function.") | |
| return parser.parse_args() | |
| def main(): | |
| # parse arguments | |
| args = parse_args() | |
| dataset = args.dataset | |
| device = args.device if args.device in DEVICES else None | |
| n_generations = args.n_generations | |
| n_runs = args.n_runs | |
| performance_weight, hardware_weight = args.performance_weight, args.hardware_weight | |
| if performance_weight + hardware_weight > 1.0 + 1e-6: | |
| error_msg = f""" | |
| Performance weight: {performance_weight}, Hardware weight: {hardware_weight} (they sum up to {performance_weight + hardware_weight}). | |
| The sum of the weights must be less than 1. | |
| """ | |
| raise ValueError(error_msg) | |
| nebulos_chunks = [] | |
| for i in range(4): # the number of chunks is 4 in this case | |
| with open(f"data/nebuloss_{i+1}.json", "r") as f: | |
| nebulos_chunks.append(json.load(f)) | |
| searchspace_dict = union_of_dicts(nebulos_chunks) | |
| # initialize the search space given dataset and device | |
| searchspace_interface = HW_NATS_FastInterface(datapath=searchspace_dict, device=args.device, dataset=args.dataset) | |
| search = GeneticSearch( | |
| searchspace=searchspace_interface, | |
| fitness_weights=np.array([performance_weight, hardware_weight]) | |
| ) | |
| # this perform the actual architecture search | |
| results = search.solve(max_generations=n_generations) | |
| print(f'{dataset}-{device.upper() if device is not None else device}') | |
| print(results[0].genotype, results[0].genotype_to_idx["/".join(results[0].genotype)], results[1]) | |
| print() | |
| if __name__=="__main__": | |
| main() | |