|
- """
- The script for testing brainscore in the openI community.
- Author: Yu Liutao @ PCL, 2021.9
- """
-
- # ---------------------------------------------------
- # Imports
- # ---------------------------------------------------
- from __future__ import print_function
- import argparse
- import torch
- from torchvision import datasets, transforms
- from torch.utils.data.dataloader import DataLoader
- import datetime
- import sys
- import os
-
- from model_tools.brain_transformation import ModelCommitment, LayerSelection, RegionLayerMap
- from brainscore import score_model
-
- # ---------------------------------------------------
- # Utils
- # ---------------------------------------------------
- benchmark_dict = {
- 'V1': 'movshon.FreemanZiemba2013public.V1-pls',
- 'V2': 'movshon.FreemanZiemba2013public.V2-pls',
- 'V4': 'dicarlo.MajajHong2015public.V4-pls',
- 'IT': 'dicarlo.MajajHong2015public.IT-pls'
- }
-
- def getBenchmark(brainRegion):
- """
- :param brainRegion:
- :return: Benchmark
- """
- return benchmark_dict.get(brainRegion, 'Invalid Brain Region')
-
-
- def find_file(directory, fileName):
- '''
- directory: the folder needed to be searched, eg /dataset
- fileName: the name of the model file, eg model.pth
- Return: the absolute path of the model file
- '''
- #import sys, os
- flag = 0
- for root, dirs, files in os.walk(directory):
- for name in files:
- if name == fileName:
- flag = 1
- filepath = os.path.join(root, name)
- print('\n Model path: {}'.format(filepath))
- return filepath
-
- if flag == 0:
- print('\n {} not found in {}'.format(fileName, directory))
- sys.exit('\n {} not found in {}'.format(fileName, directory))
-
-
- if __name__ == '__main__':
- parser = argparse.ArgumentParser(description='snn4imagenet')
- parser.add_argument('--log', default=True, type=bool, help='to print the output to log file')
- parser.add_argument('--brainRegion', default='V1', type=str, help='choose a brain region from [V1, V2, V4, IT]')
- parser.add_argument('--modelname', default='', type=str, help='the candidate model name')
- parser.add_argument('--modelpath', default='', type=str, help='path to your pretrained weights')
- parser.add_argument('--modeldescription', default='', type=str, help='one sentence less than 200 characters to describe your model briefly')
-
- args = parser.parse_args()
-
- print(parser.parse_args())
- print('\n')
-
- assert args.brainRegion != '', 'brain region (brainRegion) not provided'
- assert args.modelname != '', 'model name (modelname) not provided'
- assert args.modelpath != '', 'model path (modelpath) not provided'
- assert args.modeldescription != '', 'model description (modeldescription) not provided'
-
- # Set some key parameters (torch.device, batch_size ...)
- # device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
- os.environ['CUDA_VISIBLE_DEVICES'] = '0'
- if torch.cuda.is_available():
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
-
- # search the pretrained model file in the given directory 'modelpath'
- if args.modelname == 'alexnet':
- abs_modelpath = '/code/alexnet-owt-4df8aa71.pth'
- else:
- # abs_modelpath = find_file(args.modelpath, 'model.pth')
- abs_modelpath = args.modelpath
-
- modelsize = os.path.getsize(abs_modelpath) # bytes
- modelsize = float(modelsize/1024/1024) # MB
-
- print('\n Model name: {}'.format(args.modelname))
- print('\n Model size: {:0.1f} MB'.format(modelsize))
-
- # ---------------------------------------------------
- # log_file
- # ---------------------------------------------------
- log_file = '/code/logs/'
- try:
- # os.mkdir(log_file)
- os.makedirs(log_file) # Recursive directory creation function.
- except OSError:
- pass
- identifier = args.modelname.lower() + '.log'
- log_file += identifier
-
- if args.log:
- f = open(log_file, 'w', buffering=1)
- # f = open(log_file, 'x', buffering=1)
- else:
- f = sys.stdout
-
- # start testing
- f.write('\n Run on time: {}'.format(datetime.datetime.now()))
- f.write('\n\n Model name: {}'.format(args.modelname))
- f.write('\n Model size: {:0.1f} MB'.format(modelsize))
- f.write('\n Model description: {}'.format(args.modeldescription))
-
- # ----------------------------------------------------------------
- # Load pretrained model and import related functions/packages
- # through a custom_model.py file
- # ----------------------------------------------------------------
-
- sys.path.append("/code/")
-
- import custom_model
-
- if args.modelname == 'alexnet':
- import alexnet_eg as custom_model
-
- model_identifier = custom_model.get_model_identifier()
- layers = custom_model.get_layers(model_identifier)
- assert layers is not None
- assert isinstance(layers, list)
- assert len(layers) > 0
- model_instance = custom_model.get_model(model_identifier, abs_modelpath)
- assert model_instance is not None
-
- f.write("\n\n Load model successfully!")
- f.write("\n\n {} \n".format(model_instance))
-
- # benchmarks
- benchmark = getBenchmark(args.brainRegion)
-
- # if torch.cuda.is_available(): # and not model_instance.is_cuda:
- # model_instance.cuda(device=0)
-
- model = ModelCommitment(identifier=model_identifier, activations_model=model_instance, layers=layers)
- start_time = datetime.datetime.now()
- score = score_model(model_identifier=model_identifier, model=model, benchmark_identifier=benchmark)
- timeCost = datetime.timedelta(seconds=(datetime.datetime.now() - start_time).seconds)
-
- f.write('\n Time cost: {}'.format(timeCost))
- f.write('\n {}'.format(score)) # xarray
-
- f.write('\n\n This trial will not attend the leaderboard. Exit!')
- f.close()
- exit(0)
|