Add 2D Probabilistic UNet #119
Conversation
torch_em/model/probabilistic_unet.py
Outdated
| from torch_em.model import UNet2d | ||
| from torch_em.loss.dice import DiceLossWithLogits | ||
|
|
||
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
There was a problem hiding this comment.
Don't define this globally. Instead pass this as an argument to the functions or arguments that need it.
There was a problem hiding this comment.
I had a bit closer look and left a few more comments. In addition to addressing those we should also add a test, see for example https://github.com/constantinpape/torch-em/blob/main/test/model/test_unet.py.
torch_em/model/probabilistic_unet.py
Outdated
| num_classes=None | ||
| ): | ||
|
|
||
| super(Encoder, self).__init__() |
There was a problem hiding this comment.
This style of writing the super calls is outdated. It's better to just write super().__init__() instead.
torch_em/model/probabilistic_unet.py
Outdated
| num_classes=None | ||
| ): | ||
|
|
||
| super(AxisAlignedConvGaussian, self).__init__() |
torch_em/model/probabilistic_unet.py
Outdated
| use_tile=True | ||
| ): | ||
|
|
||
| super(Fcomb, self).__init__() |
torch_em/model/probabilistic_unet.py
Outdated
| rl_swap=False | ||
| ): | ||
|
|
||
| super(ProbabilisticUnet, self).__init__() |
torch_em/model/probabilistic_unet.py
Outdated
| """ | ||
|
|
||
| def __init__( | ||
| self, |
There was a problem hiding this comment.
Just a minor thing, but the usual style for these inits is to have only one indent:
def __init__(
self,
aa,
bb,
cc,
dd,
...
):
...
(also relevant in several places below).
torch_em/model/probabilistic_unet.py
Outdated
| return self.last_layer(output) | ||
|
|
||
|
|
||
| class ProbabilisticUnet(nn.Module): |
There was a problem hiding this comment.
Let's call it ProbabilisticUNet to be consistent with the other UNet naming.
|
Thanks for the feedback. The discussions for naming consistency (#119 (comment)), updating the indent (#119 (comment)) and style of (I'll work on adding a test for |
torch_em.model.UNet2dbackbone