Merge pull request #20 from ORNL/dev

v0.3.0

CHANGES.txt

@@ -1,2 +1,3 @@
 v0.1.0: Initial release
-v0.2.0: Spike Time Dependent Plasticity included. Heterogenous neuron API. Bug fixes to synaptic step functions.
+v0.2.0: Spike Time Dependent Plasticity included. Heterogenous neuron API. Bug fixes to synaptic step functions.
+v0.3.0: Model can be saved and loaded. Simulation can be continued with multiple calls to model.NeurmorphicModel.simulate. model.NeuromorphicModel.setup provides the option to reset the simulation temporally, while preserving weights. API functions renamed: model.NeuromorhpicModel.get_spikes -> model.NeuromorphicModel.get_spike_times and model.NeuromorphicModel.spike -> model.NeuromorphicModel.add_spike. Bug fixes to synaptic delay, compress_tensor and shared memory use in CPU mode. Customizable learning parameters.

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="SuperNeuroABM",
-    version="0.2.0",
+    version="0.3.0",
     author="Chathika Gunaratne, Prasanna Date, Shruti Kulkarni, Xi Zhang",
     author_email="[email protected]",
     packages=["superneuroabm", "superneuroabm.core"],
@@ -13,5 +13,5 @@
     long_description="""A GPU-based multi-agent simulation framework for neuromorphic computing.""",
     long_description_content_type="text/markdown",
     project_urls={"Source": "https://github.com/ORNL/superneuroabm"},
-    install_requires=["numba==0.55.1", "numpy==1.21.6", "tqdm==4.64.1"],
+    install_requires=["numba==0.55.1", "numpy==1.21.6", "tqdm==4.64.1"], 
 )

superneuroabm/core/agent.py

@@ -3,7 +3,6 @@
 from collections import OrderedDict
 from copy import copy
 import warnings
-from multiprocessing import shared_memory
 
 import numpy as np
 from numba import cuda
@@ -172,7 +171,7 @@ def set_agent_property_value(
         self._property_name_2_agent_data_tensor[property_name][
             agent_id
         ] = value
-        if dims:
+        if dims != None:
             if self._property_name_2_max_dims[property_name]:
                 if len(dims) != len(
                     self._property_name_2_max_dims[property_name]
@@ -230,16 +229,20 @@ def generate_agent_data_tensors(
             adt = convert_to_equal_side_tensor(adt, max_dims)
             if use_cuda:
                 adt = cuda.to_device(adt)
-            else:
-                # use shared memory if not cuda
-                d_size = np.dtype(dtype).itemsize * np.prod(adt.shape)
-                shm = shared_memory.SharedMemory(
-                    create=True,
-                    size=d_size,
-                    name=f"npshared{property_name}",
-                )
-                dst = np.ndarray(shape=adt.shape, dtype=dtype, buffer=shm.buf)
-                dst[:] = adt[:]
+            '''else:
+                # TODO fix shared memory usage
+                if False:
+                    d_size = np.dtype(dtype).itemsize * np.prod(adt.shape)
+                    shm = shared_memory.SharedMemory(
+                        create=True,
+                        size=d_size,
+                        name=f"npshared{property_name}",
+                    )
+                    dst = np.ndarray(
+                        shape=adt.shape, dtype=dtype, buffer=shm.buf
+                    )
+                    dst[:] = adt[:]
+            '''
             converted_agent_data_tensors.append(adt)
         return converted_agent_data_tensors
 
@@ -252,12 +255,15 @@ def update_agents_properties(
                 dt = equal_side_agent_data_tensors[i]
             else:
                 dt = equal_side_agent_data_tensors[i]
-                # Free shared memory
-                shm = shared_memory.SharedMemory(
-                    name=f"npshared{property_name}"
-                )
-                shm.close()
-                shm.unlink()
+                # TODO Free shared memory
+                '''if False:
+                    if dt.size != 0:
+                        shm = shared_memory.SharedMemory(
+                            name=f"npshared{property_name}"
+                        )
+                        shm.close()
+                        shm.unlink()
+                '''
             self._property_name_2_agent_data_tensor[
                 property_names[i]
             ] = compress_tensor(dt)

superneuroabm/core/model.py

@@ -7,9 +7,9 @@
 import heapq
 from multiprocessing import Pool, Manager
 import inspect
+import pickle
 
 from numba import cuda
-from numba.cuda.random import create_xoroshiro128p_states
 from tqdm import tqdm
 
 from superneuroabm.core.agent import AgentFactory, Breed
@@ -18,19 +18,8 @@
 class Model:
     THREADSPERBLOCK = 32
 
-    def __init__(self, name: str = "Untitled", use_cuda: bool = True) -> None:
-        self._name = name
+    def __init__(self) -> None:
         self._agent_factory = AgentFactory()
-        self._use_cuda = use_cuda
-        if self._use_cuda:
-            if not cuda.is_available():
-                raise EnvironmentError(
-                    "CUDA requested but no cuda installation detected."
-                )
-            else:
-                pass
-                # device = cuda.get_current_device()
-                # device.reset()
 
     def register_breed(self, breed: Breed) -> None:
         if self._agent_factory.num_agents > 0:
@@ -83,7 +72,20 @@ def set_agent_property_value(
     def get_agents_with(self, query: Callable) -> Set[List[Any]]:
         return self._agent_factory.get_agents_with(query=query)
 
-    def setup(self) -> None:
+    def setup(self, use_cuda: bool = True) -> None:
+        """
+        Must be called before first simulate call.
+        Initializes model and resets ticks. Readies step functions
+        and for breeds.
+
+        :param use_cuda: runs model in GPU mode.
+        """
+        self._use_cuda = use_cuda
+        if self._use_cuda:
+            if not cuda.is_available():
+                raise EnvironmentError(
+                    "CUDA requested but no cuda installation detected."
+                )
         # Create record of agent step functions by breed and priority
         self._breed_idx_2_step_func_by_priority: List[Dict[int, Callable]] = []
         heap_priority_breedidx_func = []
@@ -110,16 +112,16 @@ def setup(self) -> None:
                 )
                 last_priority = priority
 
-        # Generate agent data tensor
-        self._agent_data_tensors = (
-            self._agent_factory.generate_agent_data_tensors(self._use_cuda)
-        )
+        # Generate global data tensor
+        self._global_data_vector = [0]  # index 0 reserved for step
 
     def simulate(
         self, ticks: int, update_data_ticks: int = 1, num_cpu_proc: int = 4
     ) -> None:
-        # Generate global data tensor
-        self._global_data_vector = [0]  # index 0 reserved for step
+        # Generate agent data tensor
+        self._agent_data_tensors = (
+            self._agent_factory.generate_agent_data_tensors(self._use_cuda)
+        )
         if not self._use_cuda:
             with Pool(num_cpu_proc) as pool:
                 with Manager() as manager:
@@ -155,9 +157,10 @@ def simulate(
                         unit_divisor=1000,
                         dynamic_ncols=True,
                     ):
-                        shared_global_data_vector[0] = tick
                         for jobs_in_priority in jobs:
                             _ = list(map(smap, jobs_in_priority))
+                        shared_global_data_vector[0] += 1
+                    self._global_data_vector = list(shared_global_data_vector)
         else:
             blockspergrid = int(
                 math.ceil(
@@ -169,20 +172,39 @@ def simulate(
                 self._breed_idx_2_step_func_by_priority,
                 ticks,
             )
-            ################
-            shared_global_data_vector = [0]
             device_global_data_vector = cuda.to_device(
-                shared_global_data_vector
+                self._global_data_vector
             )
             exec(step_funcs_code_obj)
+            # Globals will have to be moved explicitly
+            self._global_data_vector = device_global_data_vector.copy_to_host()
 
         self._agent_factory.update_agents_properties(
             self._agent_data_tensors, self._use_cuda
-        )
+        )  # internals of this will automatically move data back from GPU
 
-    @property
-    def name(self) -> str:
-        return self._name
+    def save(self, app: "Model", fpath: str) -> None:
+        """
+        Saves model. Must be overridden if additional data
+        pertaining to application must be saved.
+
+        :param fpath: file path to save pickle file at
+        :param app_data: additional application data to be saved.
+        """
+        if "_agent_data_tensors" in app.__dict__:
+            del app.__dict__["_agent_data_tensors"]
+        with open(fpath, "wb") as fout:
+            pickle.dump(app, fout)
+
+    def load(self, fpath: str) -> "Model":
+        """
+        Loads model from pickle file.
+
+        :param fpath: file path to pickle file.
+        """
+        with open(fpath, "rb") as fin:
+            app = pickle.load(fin)
+        return app
 
 
 def smap(func_args):
@@ -227,8 +249,11 @@ def generate_gpu_func(
         g = cuda.cg.this_grid()
         breed_id = a0[agent_id]
         for tick in range({ticks}):
-            device_global_data_vector[0] = tick
             {sim_loop}
+            if thread_id == 0:
+                device_global_data_vector[0] += 1
+            g.sync()
+
     \nstepfunc = cuda.jit(stepfunc)
     \nstepfunc[blockspergrid, Model.THREADSPERBLOCK](
         device_global_data_vector,

superneuroabm/core/util.py

@@ -2,6 +2,7 @@
 from typing import Iterable, List, Any, Tuple, Optional
 
 import numpy as np
+from numba import cuda
 
 
 def convert_to_equal_side_tensor(
@@ -24,7 +25,7 @@ def find_max_depth(l, curr_depth=0):
         print("no dims specified... running find max depth")
         find_max_depth(tensor)
         max_dims = tuple(list(dim2maxlen.values()))
-    answer = np.full(shape=max_dims, fill_value=np.nan)
+    answer = np.full(shape=max_dims, fill_value=np.nan, dtype=np.float32)
 
     def fill_arr(arr, coord):
         if len(coord) == len(max_dims):
@@ -54,6 +55,8 @@ def compress_tensor(arr: Iterable, level: int = 0):
     else:
         new_arr = []
         for item in arr:
+            if type(item) == cuda.cudadrv.devicearray.DeviceNDArray:
+                item = item.copy_to_host()
             new_item = compress_tensor(item, level + 1)
             if (
                 (not isinstance(new_item, Iterable) and new_item != None)