Neuroballad/neuroballad_execute.py at master · FlyBrainLab/Neuroballad · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
import numpy as np
import h5py
import networkx as nx
import argparse
import itertools
import random
import pickle

import neurokernel.mpi_relaunch
import neurokernel.core_gpu as core

from neurokernel.LPU.InputProcessors.StepInputProcessor import StepInputProcessor
from neurokernel.LPU.InputProcessors.FileInputProcessor import FileInputProcessor
from neurokernel.tools.logging import setup_logger
from neurokernel.LPU.LPU import LPU

(comp_dict, conns) = LPU.lpu_parser('neuroballad_temp_model.gexf.gz')
with open('run_parameters.pickle', 'rb') as f:
    run_parameters = pickle.load(f)
with open('record_parameters.pickle', 'rb') as f:
    record_parameters = pickle.load(f)
dur = run_parameters[0]
dt = run_parameters[1]
fl_input_processor = FileInputProcessor('neuroballad_temp_model_input.h5')

from neurokernel.LPU.OutputProcessors.FileOutputProcessor import FileOutputProcessor
output_processor = FileOutputProcessor(record_parameters, 'neuroballad_temp_model_output.h5', sample_interval=1)

#Parse extra arguments
parser = argparse.ArgumentParser()
parser.add_argument('--debug', default=True,
                    dest='debug', action='store_true',
                    help='Write connectivity structures and inter-LPU routed data in debug folder')
parser.add_argument('-l', '--log', default='both', type=str,
                    help='Log output to screen [file, screen, both, or none; default:none]')
parser.add_argument('-r', '--time_sync', default=False, action='store_true',
                    help='Time data reception throughput [default: False]')
parser.add_argument('-g', '--gpu_dev', default=[0], type=int, nargs='+',
                    help='GPU device numbers [default: 0]')
parser.add_argument('-ec', '--extra_comps', default=[],
                    help='extra components to be used by LPU class')
parser.add_argument('-d', '--disconnect', default=False, action='store_true',
                    help='Run with disconnected LPUs [default: False]')
args = parser.parse_args()
file_name = None
screen = False
if args.log.lower() in ['file', 'both']:
    file_name = 'neurokernel.log'
if args.log.lower() in ['screen', 'both']:
    screen = True
logger = setup_logger(file_name=file_name, screen=screen)
man = core.Manager()

man.add(LPU, 'lpu', dt, comp_dict, conns,
        input_processors=[fl_input_processor],
        output_processors=[output_processor],
        device=args.gpu_dev[0],
        extra_comps=args.extra_comps,
        debug=True)

steps = int(dur/dt)
man.spawn()
man.start(steps=steps)
man.wait()