Tests all stable groups of neurons / nodes.
| 15 | |
| 16 | |
| 17 | class TestConnection: |
| 18 | """ |
| 19 | Tests all stable groups of neurons / nodes. |
| 20 | """ |
| 21 | |
| 22 | def __init__(self): |
| 23 | if torch.cuda.is_available(): |
| 24 | self.device = torch.device("cuda:0") |
| 25 | else: |
| 26 | self.device = torch.device("cpu:0") |
| 27 | print(f"Using device '{self.device}' for the test") |
| 28 | |
| 29 | def test_transfer(self): |
| 30 | if not torch.cuda.is_available(): |
| 31 | return |
| 32 | |
| 33 | connection_types = [ |
| 34 | Connection, |
| 35 | Conv2dConnection, |
| 36 | MaxPool2dConnection, |
| 37 | LocalConnection, |
| 38 | MeanFieldConnection, |
| 39 | SparseConnection, |
| 40 | ] |
| 41 | args = [[], [3], [3], [3, 1, 1], [], []] |
| 42 | kwargs = [{}, {}, {}, {}, {}, {"sparsity": 0.9}] |
| 43 | for conn_type, args, kwargs in zip(connection_types, args, kwargs): |
| 44 | l_a = LIFNodes(shape=[1, 28, 28]) |
| 45 | l_b = LIFNodes(shape=[1, 26, 26]) |
| 46 | connection = conn_type(l_a, l_b, *args, **kwargs) |
| 47 | |
| 48 | connection.to() |
| 49 | |
| 50 | connection_tensors = [ |
| 51 | k |
| 52 | for k, v in connection.state_dict().items() |
| 53 | if isinstance(v, torch.Tensor) and not "." in k |
| 54 | ] |
| 55 | |
| 56 | print( |
| 57 | "State dict in {} : {}".format( |
| 58 | conn_type, connection.state_dict().keys() |
| 59 | ) |
| 60 | ) |
| 61 | print("__dict__ in {} : {}".format(conn_type, connection.__dict__.keys())) |
| 62 | print("Tensors in {} : {}".format(conn_type, connection_tensors)) |
| 63 | |
| 64 | tensor_devs = [getattr(connection, k).device for k in connection_tensors] |
| 65 | print( |
| 66 | "Tensor devices {}".format(list(zip(connection_tensors, tensor_devs))) |
| 67 | ) |
| 68 | |
| 69 | for d in tensor_devs: |
| 70 | print(d, d == torch.device("cuda:0")) |
| 71 | assert d == torch.device("cuda:0") |
| 72 | |
| 73 | def test_weights(self, conn_type, shape_a, shape_b, shape_w, *args, **kwargs): |
| 74 | print("Testing:", conn_type) |