ain3007-project/monai-unet-512-no-augmentation

from monai.networks.nets import UNet          # unet model from monai (there are other models that you use with a single line)

model = UNet(
    spatial_dims=2,
    in_channels=3,
    out_channels=1,
    channels=[16, 32, 64, 128, 256, 512],  # Number of channels at each layer during contraction
    strides=(2, 2, 2, 2, 2),         # Strides for the convolutional layers
    num_res_units=4,                 # Number of residual units
    dropout=0.15,                    # Dropout rate to prevent overfitting
)