| 175 | |
| 176 | |
| 177 | class BasicUNet(nn.Module): |
| 178 | |
| 179 | def __init__( |
| 180 | self, |
| 181 | spatial_dims: int = 3, |
| 182 | in_channels: int = 1, |
| 183 | out_channels: int = 2, |
| 184 | features: Sequence[int] = (32, 32, 64, 128, 256, 32), |
| 185 | act: str | tuple = ("LeakyReLU", {"negative_slope": 0.1, "inplace": True}), |
| 186 | norm: str | tuple = ("instance", {"affine": True}), |
| 187 | bias: bool = True, |
| 188 | dropout: float | tuple = 0.0, |
| 189 | upsample: str = "deconv", |
| 190 | ): |
| 191 | """ |
| 192 | A UNet implementation with 1D/2D/3D supports. |
| 193 | |
| 194 | Based on: |
| 195 | |
| 196 | Falk et al. "U-Net – Deep Learning for Cell Counting, Detection, and |
| 197 | Morphometry". Nature Methods 16, 67–70 (2019), DOI: |
| 198 | http://dx.doi.org/10.1038/s41592-018-0261-2 |
| 199 | |
| 200 | Args: |
| 201 | spatial_dims: number of spatial dimensions. Defaults to 3 for spatial 3D inputs. |
| 202 | in_channels: number of input channels. Defaults to 1. |
| 203 | out_channels: number of output channels. Defaults to 2. |
| 204 | features: six integers as numbers of features. |
| 205 | Defaults to ``(32, 32, 64, 128, 256, 32)``, |
| 206 | |
| 207 | - the first five values correspond to the five-level encoder feature sizes. |
| 208 | - the last value corresponds to the feature size after the last upsampling. |
| 209 | |
| 210 | act: activation type and arguments. Defaults to LeakyReLU. |
| 211 | norm: feature normalization type and arguments. Defaults to instance norm. |
| 212 | bias: whether to have a bias term in convolution blocks. Defaults to True. |
| 213 | According to `Performance Tuning Guide <https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html>`_, |
| 214 | if a conv layer is directly followed by a batch norm layer, bias should be False. |
| 215 | dropout: dropout ratio. Defaults to no dropout. |
| 216 | upsample: upsampling mode, available options are |
| 217 | ``"deconv"``, ``"pixelshuffle"``, ``"nontrainable"``. |
| 218 | |
| 219 | Examples:: |
| 220 | |
| 221 | # for spatial 2D |
| 222 | >>> net = BasicUNet(spatial_dims=2, features=(64, 128, 256, 512, 1024, 128)) |
| 223 | |
| 224 | # for spatial 2D, with group norm |
| 225 | >>> net = BasicUNet(spatial_dims=2, features=(64, 128, 256, 512, 1024, 128), norm=("group", {"num_groups": 4})) |
| 226 | |
| 227 | # for spatial 3D |
| 228 | >>> net = BasicUNet(spatial_dims=3, features=(32, 32, 64, 128, 256, 32)) |
| 229 | |
| 230 | See Also |
| 231 | |
| 232 | - :py:class:`monai.networks.nets.DynUNet` |
| 233 | - :py:class:`monai.networks.nets.UNet` |
| 234 |
no outgoing calls
searching dependent graphs…