diff --git a/src/leibnetz/leibnet.py b/src/leibnetz/leibnet.py index 522f2d3..fee319d 100644 --- a/src/leibnetz/leibnet.py +++ b/src/leibnetz/leibnet.py @@ -5,6 +5,7 @@ from torch.nn import Module import numpy as np from leibnetz.nodes import Node +from funlib.learn.torch.models.conv4d import Conv4d # from model_opt.apis import optimize @@ -19,6 +20,7 @@ def __init__( nodes: Iterable, outputs: dict[str, Sequence[Tuple]], retain_buffer=True, + initialization="kaiming", ): super().__init__() full_node_list = [] @@ -35,6 +37,41 @@ def __init__( self.nodes_dict = torch.nn.ModuleDict({node.id: node for node in nodes}) self.graph = nx.DiGraph() self.assemble(outputs) + self.initialization = initialization + if initialization == "kaiming": + self.apply( + lambda m: ( + torch.nn.init.kaiming_normal_(m.weight, mode="fan_out") + if isinstance(m, torch.nn.Conv2d) + or isinstance(m, torch.nn.Conv3d) + or isinstance(m, Conv4d) + else None + ) + ) + elif initialization == "xavier": + self.apply( + lambda m: ( + torch.nn.init.xavier_normal_(m.weight) + if isinstance(m, torch.nn.Conv2d) + or isinstance(m, torch.nn.Conv3d) + or isinstance(m, Conv4d) + else None + ) + ) + elif initialization == "orthogonal": + self.apply( + lambda m: ( + torch.nn.init.orthogonal_(m.weight) + if isinstance(m, torch.nn.Conv2d) + or isinstance(m, torch.nn.Conv3d) + or isinstance(m, Conv4d) + else None + ) + ) + elif initialization is None: + pass + else: + raise ValueError(f"Unknown initialization {initialization}") self.retain_buffer = retain_buffer self.retain_buffer = True if torch.cuda.is_available(): diff --git a/src/leibnetz/nets/bio.py b/src/leibnetz/nets/bio.py new file mode 100644 index 0000000..442a46c --- /dev/null +++ b/src/leibnetz/nets/bio.py @@ -0,0 +1,167 @@ +""" +This code is taken from https://github.com/Joxis/pytorch-hebbian.git +The code is licensed under the MIT license. + +Please reference the following paper if you use this code: +@inproceedings{talloen2020pytorchhebbian, + author = {Jules Talloen and Joni Dambre and Alexander Vandesompele}, + location = {Online}, + title = {PyTorch-Hebbian: facilitating local learning in a deep learning framework}, + year = {2020}, +} +""" + +import logging +from abc import ABC, abstractmethod +import torch + +from leibnetz import LeibNet + + +class LearningRule(ABC): + + def __init__(self): + self.logger = logging.getLogger(__name__ + "." + self.__class__.__name__) + + def init_layers(self, model): + pass + + @abstractmethod + def update(self, x, w): + pass + + +class HebbsRule(LearningRule): + + def __init__(self, c=0.1): + super().__init__() + self.c = c + + def update(self, inputs, w): + # TODO: Needs re-implementation + d_ws = torch.zeros(inputs.size(0)) + for idx, x in enumerate(inputs): + y = torch.dot(w, x) + + d_w = torch.zeros(w.shape) + for i in range(y.shape[0]): + for j in range(x.shape[0]): + d_w[i, j] = self.c * x[j] * y[i] + + d_ws[idx] = d_w + + return torch.mean(d_ws, dim=0) + + +class KrotovsRule(LearningRule): + """Krotov-Hopfield Hebbian learning rule fast implementation. + + Original source: https://github.com/DimaKrotov/Biological_Learning + + Args: + precision: Numerical precision of the weight updates. + delta: Anti-hebbian learning strength. + norm: Lebesgue norm of the weights. + k: Ranking parameter + """ + + def __init__(self, precision=1e-30, delta=0.4, norm=2, k=2, normalize=False): + super().__init__() + self.precision = precision + self.delta = delta + self.norm = norm + self.k = k + self.normalize = normalize + + def init_layers(self, layers: list): + for layer in [lyr.layer for lyr in layers]: + if type(layer) == torch.nn.Linear or type(layer) == torch.nn.Conv2d: + layer.weight.data.normal_(mean=0.0, std=1.0) + + def update(self, inputs: torch.Tensor, weights: torch.Tensor): + batch_size = inputs.shape[0] + num_hidden_units = weights.shape[0] + input_size = inputs[0].shape[0] + assert ( + self.k <= num_hidden_units + ), "The amount of hidden units should be larger or equal to k!" + + # TODO: WIP + if self.normalize: + norm = torch.norm(inputs, dim=1) + norm[norm == 0] = 1 + inputs = torch.div(inputs, norm.view(-1, 1)) + + inputs = torch.t(inputs) + + # Calculate overlap for each hidden unit and input sample + tot_input = torch.matmul( + torch.sign(weights) * torch.abs(weights) ** (self.norm - 1), inputs + ) + + # Get the top k activations for each input sample (hidden units ranked per input sample) + _, indices = torch.topk(tot_input, k=self.k, dim=0) + + # Apply the activation function for each input sample + activations = torch.zeros((num_hidden_units, batch_size)) + activations[indices[0], torch.arange(batch_size)] = 1.0 + activations[indices[self.k - 1], torch.arange(batch_size)] = -self.delta + + # Sum the activations for each hidden unit, the batch dimension is removed here + xx = torch.sum(torch.mul(activations, tot_input), 1) + + # Apply the actual learning rule, from here on the tensor has the same dimension as the weights + norm_factor = torch.mul( + xx.view(xx.shape[0], 1).repeat((1, input_size)), weights + ) + ds = torch.matmul(activations, torch.t(inputs)) - norm_factor + + # Normalize the weight updates so that the largest update is 1 (which is then multiplied by the learning rate) + nc = torch.max(torch.abs(ds)) + if nc < self.precision: + nc = self.precision + d_w = torch.true_divide(ds, nc) + + return d_w + + +class OjasRule(LearningRule): + + def __init__(self, c=0.1): + super().__init__() + self.c = c + + def update(self, inputs, w): + # TODO: needs re-implementation + d_ws = torch.zeros(inputs.size(0), *w.shape) + for idx, x in enumerate(inputs): + y = torch.mm(w, x.unsqueeze(1)) + + d_w = torch.zeros(w.shape) + for i in range(y.shape[0]): + for j in range(x.shape[0]): + d_w[i, j] = self.c * y[i] * (x[j] - y[i] * w[i, j]) + + d_ws[idx] = d_w + + return torch.mean(d_ws, dim=0) + + +def convert_to_bio(model: LeibNet, learning_rule: LearningRule, **kwargs): + """Converts a LeibNet model to use local bio-inspired learning rules. + + Args: + model (LeibNet): Initial LeibNet model to convert. + learning_rule (LearningRule): Learning rule to apply to the model. Can be `HebbsRule`, `KrotovsRule` or `OjasRule`. + + Returns: + _type_: _description_ + """ + + def hook(module, args, kwargs, output): ... + + for module in model.modules(): + if len(module._parameters) > 0: + module.register_forward_hook(hook, with_kwargs=True) + + return model