defineUNet3p: defineUnet3p

In geodl: Geospatial Semantic Segmentation with Torch and Terra

defineUnet3p

Description

Define a UNet3+ architecture for use in luz training loop.

Usage

defineUNet3p(
  inChn = 3,
  nCls = 2,
  enChn = c(16, 32, 64, 128),
  outChn = 64,
  btnChn = 256,
  useASPP = TRUE,
  dilRates = c(6, 12, 18),
  dilChn = c(256, 256, 256, 256),
  negative_slope = 0.01,
  useDS = FALSE
)

Arguments

inChn	Number of channels, bands, or predictor variables in the input image or raster data. Default is 3.
nCls	Number of classes being differentiated. For a binary classification, this can be either 1 or 2. If 2, the problem is treated as a multiclass problem, and a multiclass loss metric should be used. Default is 3. double convolution operation. Default is FALSE or the ASPP module is not used as the bottleneck.
enChn	Vector of 4 integers defining the number of output feature maps for each of the four encoder blocks. Default is 16, 32, 64, and 128. maps for each of the 4 decoder blocks. Default is 128, 64, 32, and 16.
outChn	Number of output channels for each decoder block. Default is 64.
btnChn	Number of output feature maps from the bottleneck block. Default is 256.
useASPP	TRUE or FALSE. Whether to use an ASPP module as the bottleneck as opposed to a double convolution operation. Default is FALSE or the ASPP module is not used as the bottleneck.
dilRates	Vector of 3 values specifying the dilation rates used in the ASPP module. Default is 6, 12, and 18.
dilChn	Vector of 4 values specifying the number of channels to produce at each dilation rate within the ASPP module. Default is 256 for each dilation rate.
negative_slope	Specifies the negative slope term for leaky ReLU activation. Default is 0.01.
useDS	TRUE or FALSE. Whether or not to use deep supervision. If TRUE, four predictions are made, one at each decoder block resolution, and the predictions are returned as a list object containing the 4 predictions. If FALSE, only the final prediction at the original resolution is returned. Default is FALSE or deep supervision is not implemented.

Details

Define a UNet3+-like architecture for use in luz training loop. User can specify the number of output feature maps from each encoder and decoder block and the bottleneck block. Deep supervision can also be implemented. The default bottleneck block can be replaced with a atrous spatial pyramid pooling module. Leaky ReLU is used throughout. A variable number of input predictor variables and output classes can be defined.

The architecture was inspired by:

Huang, H., Lin, L., Tong, R., Hu, H., Zhang, Q., Iwamoto, Y., Han, X., Chen, Y.W. and Wu, J., 2020, May. Unet 3+: A full-scale connected unet for medical image segmentation. In ICASSP 2020-2020 IEEE international conference on acoustics, speech and signal processing (ICASSP) (pp. 1055-1059). IEEE.

Value

UNet3+ model using nn_module().

Examples

library(torch)
model <- defineUNet3p(inChn=4,
nCls=2,
enChn = c(16,32,64,128),
outChn = 64,
btnChn = 256,
useASPP = TRUE,
dilRates=c(6,12,18),
dilChn=c(256,256,256,256),
negative_slope=0.01,
useDS = FALSE)$to(device="cuda")
t1 <- torch::torch_rand(2,4,512,512)$to(device="cuda")
tp <- model(t1)

geodl documentation built on Nov. 12, 2025, 5:07 p.m.