search

fepegar / unet

"pip install unet": PyTorch Implementation of 1D, 2D and 3D U-Net architecture.
MIT License
148 stars 22 forks source link

3d implementation for a regression problem. #32

Closed jakemanger closed 2 years ago

jakemanger commented 2 years ago

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?

fepegar commented 2 years ago

Hi, @jakemanger. Sorry I'm late, I just noticed today that I wasn't watching the repository.

For regression, I would remove the classifier. In your case:

self._model.classifier = torch.nn.Identity()