README.md

topoloss

Induce brain-like topographic structure in your neural networks.

Read the paper (ICLR 2025), check out the colab notebook and play with the pre-trained models 🤗

pip install topoloss

Example

import torchvision.models as models
from topoloss import TopoLoss, LaplacianPyramid

model = models.resnet18(weights = "DEFAULT")

topo_loss = TopoLoss(
    losses = [
        LaplacianPyramid.from_layer(
            model=model,
            layer = model.fc, ## supports nn.Linear and nn.Conv2d
            factor_h=8.0, 
            factor_w=8.0, 
            scale = 1.0
        ),
    ],
)
loss = topo_loss.compute(model=model)
## >>> tensor(0.8407, grad_fn=<DivBackward0>)
loss.backward()

loss_dict = topo_loss.compute(model=model, reduce_mean = False) ## {"fc": }
## >>> {'fc': tensor(0.8407, grad_fn=<MulBackward0>)}

Running tests

pytest -vvx tests