Home

Awesome

Python 3.9 GitHub

πŸ”₯ OctoPyTorch: Segmentation Neural Networks πŸ”₯

No backbones needed, with a focus on medical images

Implementation of the segmentation neural networks for PyTorch with new features such as:

For the Tiramisu architecture:

Roadmap

Support for the following neural networks:

Getting Started

The package can be installed from the repository with:

> pip3 install octopytorch

You can try the model in Python with:

from functools import partial
import torch
from torch import nn
import octopytorch as octo

module_bank = octo.DEFAULT_MODULE_BANK.copy()
# Dropout
module_bank[octo.ModuleType.DROPOUT] = partial(nn.Dropout2d, p=0.2, inplace=True)
# Every activation in the model is going to be a GELU (Gaussian Error Linear 
# Units function). GELU(x) = x * Ξ¦(x)
# See: https://pytorch.org/docs/stable/generated/torch.nn.GELU.html
module_bank[octo.ModuleType.ACTIVATION] = nn.GELU
# Example for segmentation:
module_bank[octo.ModuleType.ACTIVATION_FINAL] = partial(nn.LogSoftmax, dim=1)
# Example for regression (default):
#module_bank[octo.ModuleType.ACTIVATION_FINAL] = nn.Identity

model = octo.models.Tiramisu(
    in_channels = 3,          # RGB images
    out_channels = 5,         # 5-channel output (5 classes)
    init_conv_filters = 48,   # Number of channels outputted by the 1st convolution
    structure = (
        [4, 4, 4, 4, 4],      # Down blocks
        4,                    # bottleneck layers
        [4, 4, 4, 4, 4],      # Up blocks
    ),
    growth_rate = 12,         # Growth rate of the DenseLayers
    compression = 1.0,        # No compression
    early_transition = False, # No early transition
    include_top = True,       # Includes last layer and activation
    checkpoint = False,       # No memory checkpointing
    module_bank = module_bank # Modules to use
)

# Initializes all the convolutional kernel weights.
model.initialize_kernels(nn.init.kaiming_uniform_, conv=True)
# Shows some information about the model.
model.summary()

This example tiramisu network has a depth of len(down_blocks) = 5, meaning that the input images should be at least 32x32 pixels (i.e. 2^5=32).

Documentation

The parameters of the constructor are explained as following:

Module bank

The Tiramisu base layers (e.g. Conv2D, activation functions, etc.) can be set to different types of layers. This was introduced to wrap many arguments of the main class under the same object and increase the flexibility to change layers.

The layers that can be redefined are:

Notes:

Tips and tricks

Built With

Contributing

See also the list of contributors who participated in this project. For contributing, make sure the code passes the checks of Pylama, Bandit and Mypy. Additionally, the code is formatted with Black.

License

This project is licensed under the MIT License - see the LICENSE.md file for details.

Acknowledgments

Many thanks to @RaphaelaHeil for her much appreciated advices on best practices.