fepegar / unet

"pip install unet": PyTorch Implementation of 1D, 2D and 3D U-Net architecture.

U-Net

PyTorch implementation of 2D and 3D U-Net.

The U-Net architecture was first described in Ronneberger et al. 2015, U-Net: Convolutional Networks for Biomedical Image Segmentation. The 3D version was described in Çiçek et al. 2016, 3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation.

Installation

pip install unet

Credits

If you used this code for your research, please cite this repository using the information available on its Zenodo entry:

Pérez-García, Fernando. (2020). fepegar/unet: PyTorch implementation of 2D and 3D U-Net (v0.7.5). Zenodo. https://doi.org/10.5281/zenodo.3697931

unet's Issues

Is the 1D implementation still available?

Is the 1D implementation still available?

Thanks!

Monte carlo dropout may be misleading

unet/unet/unet.py

Line 32 in 9f64483

monte_carlo_dropout: float = 0,

Monte carlo dropout refers more generally to dropout used during test time, in addition to train time, in order to produce a distribution of outputs whose uncertainty can be measured. The monte carlo dropout argument implemented here is a specific implementation from the paper "On the Compactness, Efficiency, and Representation of 3D Convolutional Networks: Brain Parcellation as a Pretext Task". It may be better to rename as "monte carlo dropout layer" to reflect its use.

3D Unet size mismatch

• unet version: latest version
• Python version: 3.6.9
• Operating System: Linux

Description

I am using your 3D unet to train on images of size (256, 256, 256)

What I Did

I am using MONAI and pytorch for this. Since MONAI's UNet is not completely correct, I've decided to test yours.
However, I am getting the error.
Sizes of tensor must match except in dimension 2. Got 48 and 49 (The offending index is 0).

I am assuming the input as batch size and then channels, so (1, 1, 256, 256, 256).

But what is going wrong here?

Thank you.

3d implementation for a regression problem.

Thank you for this great U-net implementation.
This isn't an issue, but I couldn't find a discussion page, so thought I should post this here. I hope that is OK.

My application requires a single input volume and a single output volume for a regression task. Some other U-net implementations I have seen have a segmentation option, which in this case I would assume should be false.

In the absence of this, would creating a U-net for regression with the following code be the correct usage of this implementation for regression?

 self._model = UNet3D(
            in_channels=1,
            out_classes=1,
            num_encoding_blocks=5,
            out_channels_first_layer=64,
            normalization='batch',
            pooling_type='avg',
            upsampling_type='linear',
            padding=True,
            activation='ReLU',
            dimensions=3,
            dropout=0.25,
        )

Is any modification to the classifier: ConvolutionalBlock in unet.py required?

How to cite

Hi,

Could you explain how to cite if I have used this code. I am unsure how to use the information available in the Zenodo entry as suggested by the readme.

Thanks,
Pritesh

Initial Update

The bot created this issue to inform you that pyup.io has been set up on this repo.
Once you have closed it, the bot will open pull requests for updates as soon as they are available.

Documentation offline

Description

Hi. I just ran across your implementation of unet and wanted to read the documentation, but the link leads to a blank readthedocs page.

Have a nice day!

Does this work on Windows?

Does this work on Windows, specifically Windows 10? Asking because I downloaded another someone else's implementation of 3D Unet on GitHub only to realize later that it doesn't work on Windows.

Not all input dimensions work out

Please document how to train and predict, particularly regarding input dimensions.

Minimal failing example

import torch
from unet import UNet2D

device = 'cuda:0'
spatial_dimension = 511
n_samples = 2
n_channels = 1

X = torch.zeros((n_samples, n_channels, spatial_dimension, spatial_dimension),
                dtype=torch.float32, device=device)

model = UNet2D(in_channels=n_channels, residual=True).to(device)
model.forward(X)

Error:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-27-cb3f6be760ae> in <module>
     11 
     12 model = UNet2D(in_channels=n_channels, residual=True).to(device)
---> 13 model.forward(X)

unet/unet.py in forward(self, x)
    120         skip_connections, encoding = self.encoder(x)
    121         encoding = self.bottom_block(encoding)
--> 122         x = self.decoder(skip_connections, encoding)
    123         if self.monte_carlo_layer is not None:
    124             x = self.monte_carlo_layer(x)

torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

unet/decoding.py in forward(self, skip_connections, x)
     59         zipped = zip(reversed(skip_connections), self.decoding_blocks)
     60         for skip_connection, decoding_block in zipped:
---> 61             x = decoding_block(skip_connection, x)
     62         return x
     63 

torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

unet/decoding.py in forward(self, skip_connection, x)
    129         x = self.upsample(x)
    130         skip_connection = self.center_crop(skip_connection, x)
--> 131         x = torch.cat((skip_connection, x), dim=CHANNELS_DIMENSION)
    132         if self.residual:
    133             connection = self.conv_residual(x)

RuntimeError: Sizes of tensors must match except in dimension 2. Got 62 and 63 (The offending index is 0)

Different kernel sizes for different dimensions

Can you maybe add this feature, please?

I think the main change would be

• From

class ConvolutionalBlock(nn.Module):
    def __init__(
            self,
            dimensions: int,
            in_channels: int,
            out_channels: int,
            normalization: Optional[str] = None,
            kernel_size: int = 3,
            activation: Optional[str] = 'ReLU',
            preactivation: bool = False,
            padding: int = 0,
            padding_mode: str = 'zeros',
            dilation: Optional[int] = None,
            dropout: float = 0,
            ):

• To

class ConvolutionalBlock(nn.Module):
    def __init__(
            self,
            dimensions: int,
            in_channels: int,
            out_channels: int,
            normalization: Optional[str] = None,
            kernel_size: int = Union[int, Sequence[int]],
            activation: Optional[str] = 'ReLU',
            preactivation: bool = False,
            padding: int = 0,
            padding_mode: str = 'zeros',
            dilation: Optional[int] = None,
            dropout: float = 0,
            ):

I can create a PR if you prefer this way.

Warning about floor division

From TorchIO tutorials:

/usr/local/lib/python3.7/dist-packages/unet/decoding.py:146: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').
  half_crop = crop // 2

Error with inplace operations due to recent changes in pytorch

Hi, running the following script with torch=2.0.0:

import unet
import torch
loss_fn = torch.nn.BCEWithLogitsLoss()
input = torch.ones([8, 3, 24, 24])
targets = torch.ones([8, 10, 24, 24])
unet_model = unet.UNet2D(in_channels=3, out_classes=10, residual=True, num_encoding_blocks=2)
out = unet_model(input)
loss = loss_fn(targets, out)
loss.backward()

gave the following error:

---

RuntimeError Traceback (most recent call last)
/tmp/ipykernel_5969/4012461463.py in
5 out = unet_model(input)
6 loss = loss_fn(targets, out)
----> 7 loss.backward()

~/.local/lib/python3.10/site-packages/torch/_tensor.py in backward(self, gradient, retain_graph, create_graph, inputs)
485 inputs=inputs,
486 )
--> 487 torch.autograd.backward(
488 self, gradient, retain_graph, create_graph, inputs=inputs
489 )

~/.local/lib/python3.10/site-packages/torch/autograd/init.py in backward(tensors, grad_tensors, retain_graph, create_graph,
grad_variables, inputs)
198 # some Python versions print out the first line of a multi-line function
199 # calls in the traceback and some print out the last line
--> 200 Variable.execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
201 tensors, grad_tensors, retain_graph, create_graph, inputs,
202 allow_unreachable=True, accumulate_grad=True) # Calls into the C++ engine to run the backward pass

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [8, 64, 24, 24]], which is output 0 of ReluBackward0, is at version 1; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

Torch does not allow the += operator on tensors anymore.
To fix it, you need to change two lines of code that are considered inplace operations in the newer torch version:
decoding.py l.136: x += connection => x = x + connection
encoding.py l.146: x += connection => x = x + connection
I made a clone of this rep and fixed the error locally this way, but I did not have the permission to push.
Best Regards,
Johannes Dollinger

Forward pass going twice through ConvolutionalBlock layers

• unet version: 0.7.3
• Python version: 3.7.4
• Operating System: Linux Mint 19.2

Description

Working on a function to generate a Keras-like summary for Pytoch models from pytorch-summary I came across an issue with the ConvolutionalBlock class: layers are registered twice in the model leading to weird consequences.

Indeed, in a forward pass, the model goes twice through every convolutional layer in a ConvolutionalBlock. Here is a code to illustrate this behaviour:

import torch
from unet import UNet3D

def hook(module, i, o):
    class_name = str(module.__class__).split(".")[-1].split("'")[0]
    if class_name != 'Conv3d':
        return
    print(module.id)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
input = torch.zeros((1, 3, 64, 64, 64)).to(device)
model = UNet3D(3, 3, num_encoding_blocks=2).to(device)
model.eval()

hooks = []
model.apply(lambda module: hooks.append(module.register_forward_hook(hook)))
model(input)
for h in hooks:
    h.remove()

To run it, you need to add this small change to the ConvolutionalBlock class in order to assign a unique id to every convolutional layer:

conv_layer.id = os.urandom(24).hex()

The above snippet produces the following output where you can see that every id appears twice:

Fix

A simple fix seems to be to remove the following lines:

self.conv_layer = conv_layer
self.norm_layer = norm_layer
self.activation_layer = activation_layer
self.dropout_layer = dropout_layer

and add the out_channels variable as an attribute of the ConvolutionalBlock class to retrieve it in the EncodingBlock class through self.conv2.out_channels instead of self.conv2.conv_layer.out_channels.

If this fix seems reasonable I can take care of the PR.