U-Net | Flow Arena

Variants

UNet3D

20.8M

3D volumetric version using 3D convolutions, suitable for all 3D reservoir datasets (Arena, CO2 Nested, 3D Channels).

Conv: Conv3d

Norm: GroupNorm

Exact params: 20,827,592

Benchmarks:

Flow Arena Two-Phase Oil-Water CO2 Nested Two-Phase 3D Channels CO2 Radial

UNet2D

7.3M

2D version for Cartesian grids, used on Two-Phase Oil-Water and CO2 Radial datasets.

Conv: Conv2d

Norm: GroupNorm

Exact params: 7,300,752

Benchmarks:

Two-Phase Oil-Water CO2 Radial

Training Provenance

These values are read from the generated benchmark result metadata, not from static method-page copy. When a model is retrained and the website data is regenerated, this section updates with the result records.

Arena training recipe

Epochs: 200

Batch size: 32

Learning rate: 0.0001

Loss: MSEGradientLoss, CombinedGradientLoss

Arena feature sets

Flow Arena input channels vary by fault complexity. Current result records for this method include 3 feature-set definitions.

PORO PERMX SWAT_0 PRESSURE_0 MULTX MULTY DEPTH

Strengths & Weaknesses

Strengths

Best-performing architecture on arena benchmarks and consistently strong across all dataset families
Skip connections preserve fine spatial details needed for sharp saturation fronts
Residual blocks with GroupNorm enable stable training even with small batch sizes
Relatively fast training convergence compared to transformer-based architectures
Well-understood architecture with extensive literature on tuning and modifications

Weaknesses

Purely local receptive field — may miss long-range pressure communication across the reservoir
Fixed spatial resolution hierarchy — cannot adapt to multi-scale features without architectural changes
Parameter count scales steeply with 3D convolutions (20.8M for 3D vs 7.3M for 2D)
No explicit mechanism for learning physical symmetries or conservation laws

Results Across Benchmarks

All results for U-Net variants across all benchmark families. Sorted by the first field's rel-L² error by default.

Flow Arena

Property

Fault

Control

Loss

28 results

					Sat (frac)			Pres (bar)
#	Dataset	Model	Loss	Epochs	rel-L²▲	MRE	MAE	rel-L²	MRE	MAE	Params
1	geostat_nf_rate	UNet3D	AbsLp(p=2)	200	0.1436	0.2628	0.038677	0.0302	0.0260	5.741697	20.8M
2	geostat_nf_bhp	UNet3D	AbsLp(p=2)	200	0.1470	0.2833	0.036387	0.0219	0.0184	4.658811	20.8M
3	geostat_nf_bhp	UNet3D	MSE	200	0.1478	0.2761	0.035909	0.0228	0.0192	4.851583	20.8M
4	geostat_nf_rate	UNet3D	MSE	200	0.1555	0.3080	0.039246	0.0313	0.0264	5.805879	20.8M
5	channels_nf_bhp	UNet3D	AbsLp(p=2)	200	0.1722	0.3236	0.048437	0.0267	0.0220	5.438548	20.8M
6	channels_nf_bhp	UNet3D	MSE	200	0.1766	0.3465	0.048261	0.0266	0.0217	5.393267	20.8M
7	geostat_zt_rate	UNet3D	AbsLp(p=2)	200	0.1848	0.3381	0.046976	0.0456	0.0360	8.486194	20.8M
8	geostat_zt_rate	UNet3D	MSE	200	0.1854	0.3418	0.046111	0.0469	0.0365	8.569011	20.8M
9	geostat_vt_rate	UNet3D	MSE	200	0.1894	0.3519	0.054428	0.0461	0.0359	8.486502	35.9M
10	geostat_zt_bhp	UNet3D	MSE	200	0.1894	0.3368	0.044286	0.0307	0.0248	6.210294	20.8M
11	geostat_zt_bhp	UNet3D	AbsLp(p=2)	200	0.1901	0.3473	0.045846	0.0307	0.0251	6.277672	20.8M
12	geostat_vt_bhp	UNet3D	MSE	200	0.1942	0.3673	0.053038	0.0295	0.0239	5.988580	35.9M
13	geostat_vt_rate	UNet3D	AbsLp(p=2)	200	0.2010	0.3629	0.055984	0.0467	0.0364	8.458871	35.9M
14	geostat_vt_bhp	UNet3D	AbsLp(p=2)	200	0.2022	0.3808	0.053607	0.0303	0.0243	6.077247	35.9M
15	channels_nf_rate	UNet3D	MSE	200	0.2089	0.3816	0.052728	0.0437	0.0348	7.558574	20.8M
16	channels_nf_rate	UNet3D	AbsLp(p=2)	200	0.2104	0.3861	0.054513	0.0449	0.0354	7.782996	20.8M
17	geostat_vt_rate	UNet3D	MSE+grad(w=1)	200	0.2109	0.3744	0.054545	0.0481	0.0372	8.508275	35.9M
18	geostat_vt_bhp	UNet3D	MSE+grad(w=1)	200	0.2134	0.3590	0.052282	0.0341	0.0273	6.820544	35.9M
19	channels_zt_bhp	UNet3D	MSE	200	0.2230	0.4164	0.054551	0.0319	0.0249	6.217045	20.8M
20	channels_zt_bhp	UNet3D	AbsLp(p=2)	200	0.2232	0.4314	0.055996	0.0323	0.0259	6.439610	20.8M
21	channels_vt_bhp	UNet3D	MSE	200	0.2384	0.4567	0.065744	0.0344	0.0266	6.697660	35.9M
22	channels_vt_bhp	UNet3D	MSE+grad(w=1)	200	0.2402	0.4734	0.067157	0.0344	0.0269	6.753836	35.9M
23	channels_zt_rate	UNet3D	AbsLp(p=2)	200	0.2457	0.4570	0.062272	0.0537	0.0413	9.744350	20.8M
24	channels_zt_rate	UNet3D	MSE	200	0.2457	0.4466	0.059286	0.0549	0.0424	10.070366	20.8M
25	channels_vt_bhp	UNet3D	AbsLp(p=2)	200	0.2506	0.5251	0.069934	0.0354	0.0276	6.904829	35.9M
26	channels_vt_rate	UNet3D	MSE+grad(w=1)	200	0.2574	0.4750	0.071365	0.0540	0.0409	9.889451	35.9M
27	channels_vt_rate	UNet3D	MSE	200	0.2596	0.5028	0.073584	0.0552	0.0418	10.178999	35.9M
28	channels_vt_rate	UNet3D	AbsLp(p=2)	200	0.2665	0.5476	0.076183	0.0564	0.0434	10.478560	35.9M

Two-Phase Oil-Water

Loss

16 results

					Pres (bar)			Sat (frac)
#	Target	Model	Loss	Epochs	rel-L²▲	MRE	MAE	rel-L²	MRE	MAE	Params
1	combined	UNet3D	BadawiCombined	2000	0.0103	0.0086	1.430230	0.0182	0.0140	0.004765	20.8M
2	combined	UNet3D	BadawiCombined	200	0.0142	0.0124	2.176428	0.0254	0.0207	0.007156	20.8M
3	combined	UNet2D	BadawiCombined	4000	0.0143	0.0121	2.038741	0.0202	0.0151	0.005132	29.1M
4	combined	UNet2D	BadawiCombined	4000	0.0150	0.0128	2.136473	0.0200	0.0147	0.005032	16.4M
5	combined	UNet2D	BadawiCombined	4000	0.0165	0.0142	2.364149	0.0209	0.0154	0.005240	29.1M
6	combined	UNet2D	BadawiCombined	500	0.0179	0.0149	2.469467	0.0239	0.0180	0.006175	16.4M
7	pressure	UNet2D	BadawiSingleField	500	0.0215	0.0192	3.150069	-	-	-	16.4M
8	combined	UNet2D	RelLp	200	0.0243	0.0204	3.415011	0.0268	0.0201	0.006798	7.3M
9	pressure	UNet3D	BadawiSingleField	200	0.0284	0.0268	4.242195	-	-	-	20.8M
10	combined	UNet2D	RelLp	200	0.0284	0.0259	4.494401	0.0243	0.0187	0.006455	16.4M
11	pressure	UNet2D	RelLp	200	0.0317	0.0278	4.753270	-	-	-	7.3M
12	pressure	UNet2D	RelLp	200	0.0351	0.0316	5.364382	-	-	-	16.4M
13	saturation	UNet2D	RelLp	200	-	-	-	0.0319	0.0243	0.008113	7.3M
14	saturation	UNet2D	RelLp	200	-	-	-	0.0299	0.0227	0.007628	16.4M
15	saturation	UNet2D	BadawiSingleField	500	-	-	-	0.0302	0.0234	0.007880	16.4M
16	saturation	UNet3D	BadawiSingleField	200	-	-	-	0.0240	0.0181	0.006264	20.8M

CO2 Radial

Loss

23 results

					Pres (bar)			Sat (frac)
#	Target	Model	Loss	Epochs	rel-L²▲	MRE	MAE	rel-L²	MRE	MAE	Params
1	pressure	UNet3D	UFNODerivLoss	500	0.0030	0.0008	0.181463	-	-	-	20.8M
2	pressure	UNet2D	RelLp	500	0.0031	0.0008	0.193750	-	-	-	16.3M
3	pressure	UNet2D	RelLp	500	0.0032	0.0010	0.217995	-	-	-	29.0M
4	combined	UNet2D	RelLp	500	0.0034	0.0010	0.232792	0.0708	5.9654	0.001713	16.3M
5	pressure	UNet2D	RelLp	500	0.0037	0.0010	0.235549	-	-	-	29.2M
6	pressure	UNet2D	UFNODerivLoss	500	0.0037	0.0011	0.248648	-	-	-	29.0M
7	pressure	UNet2D	UFNODerivLoss	500	0.0039	0.0011	0.250622	-	-	-	16.3M
8	pressure	UNet2D	UFNODerivLoss	500	0.0039	0.0011	0.246458	-	-	-	16.3M
9	pressure	UNet3D	RelLp	500	0.0040	0.0013	0.286640	-	-	-	20.8M
10	pressure	UNet3D	RelLp	500	0.0042	0.0014	0.291506	-	-	-	20.8M
11	pressure	UNet3D	UFNODerivLoss	500	0.0045	0.0014	0.296649	-	-	-	20.8M
12	saturation	UNet2D	RelLp	500	-	-	-	0.0653	6.6551	0.001726	16.3M
13	saturation	UNet2D	RelLp	500	-	-	-	0.0811	8.2827	0.002471	29.2M
14	saturation	UNet2D	RelLp	500	-	-	-	0.0811	8.2827	0.002471	29.2M
15	saturation	UNet2D	UFNODerivLoss	500	-	-	-	0.0829	6.2847	0.002182	16.3M
16	saturation	UNet2D	UFNODerivLoss	500	-	-	-	0.0804	6.4845	0.002141	16.3M
17	saturation	UNet2D	RelLp	500	-	-	-	0.0703	5.2798	0.001675	29.0M
18	saturation	UNet2D	RelLp	500	-	-	-	0.0703	5.2798	0.001675	29.0M
19	saturation	UNet2D	UFNODerivLoss	500	-	-	-	0.0705	6.5478	0.001767	29.0M
20	saturation	UNet3D	RelLp	500	-	-	-	0.0657	6.1184	0.001715	20.8M
21	saturation	UNet3D	RelLp	500	-	-	-	0.1005	5.6692	0.001667	20.8M
22	saturation	UNet3D	UFNODerivLoss	500	-	-	-	0.0954	6.2312	0.002083	20.8M
23	saturation	UNet3D	UFNODerivLoss	500	-	-	-	0.0634	4.0132	0.001151	20.8M

CO2 Nested

3 results

				Pres (bar)
#	Model	Loss	Epochs	rel-L²▲	MRE	MAE	Params
1	UNet3D	RelLp	800	4.12e-5	8.46e-6	0.001864	20.8M
2	UNet3D	RelLp	500	4.25e-5	8.99e-6	0.001982	20.8M
3	UNet3D	RelLp	200	5.55e-5	1.40e-5	0.003083	20.8M

Two-Phase 3D Channels

6 results

					Pres (bar)			Sat (frac)
#	Target	Model	Loss	Epochs	rel-L²▲	MRE	MAE	rel-L²	MRE	MAE	Params
1	pressure	UNet3D	RelLp	500	0.0007	0.0005	0.163858	-	-	-	20.8M
2	pressure	UNet3D	RelLp	200	0.0007	0.0006	0.181097	-	-	-	20.8M
3	combined	UNet3D	RelLp	500	0.0009	0.0006	0.206761	0.1278	0.0758	0.015825	20.8M
4	combined	UNet3D	RelLp	200	0.0010	0.0007	0.231114	0.1320	0.0899	0.018015	20.8M
5	saturation	UNet3D	RelLp	200	-	-	-	0.1323	0.0807	0.017179	20.8M
6	saturation	UNet3D	RelLp	500	-	-	-	0.1298	0.0709	0.015528	20.8M

Training Curves

UNet3D - Flow Arena (channels_nf_bhp)

UNet3D - Flow Arena (channels_nf_rate)

UNet3D - Flow Arena (channels_vt_bhp)

UNet3D - Flow Arena (channels_vt_rate)

UNet3D - Flow Arena (channels_zt_bhp)

UNet3D - Flow Arena (channels_zt_rate)

UNet3D - Flow Arena (geostat_nf_bhp)

UNet3D - Flow Arena (geostat_nf_rate)

Architecture

U-Net was originally introduced by Ronneberger et al. (2015) for biomedical image segmentation and has since become the de facto standard architecture for dense prediction tasks. Its encoder-decoder structure with skip connections allows it to combine high-level semantic features from the bottleneck with fine-grained spatial details from the encoder — critical for accurately predicting spatially varying fields like pressure and saturation.

Our implementation uses residual blocks (ResNet-style) with GroupNorm normalization instead of BatchNorm, which provides more stable training with small batch sizes common in 3D volumetric applications. The architecture follows a standard 4-level hierarchy with channel doubling at each downsampling stage.

For 3D datasets, we use volumetric 3D convolutions throughout (UNet3D). For 2D datasets, we use standard 2D convolutions (UNet2D). Both variants share the same architectural design — residual blocks, skip connections, and GroupNorm — differing only in the spatial dimensionality of the convolution kernels.

U-Net serves as our primary baseline architecture. Its consistent performance across all benchmarks makes it the reference point against which more specialized architectures are measured.

Model Summary

UNet3D

PaddedUNet3D(
  (conv_in): Conv3d(5, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
  (down_blocks): ModuleList(
    (0): DownBlock3D(
      (resnets): ModuleList(
        (0): ResnetBlock3D(
          (conv1): Conv3d(32, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm1): GroupNorm(2, 64, eps=1e-05, affine=True)
          (conv2): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm2): GroupNorm(2, 64, eps=1e-05, affine=True)
          (shortcut): Conv3d(32, 64, kernel_size=(1, 1, 1), stride=(1, 1, 1))
        )
        (1): ResnetBlock3D(
          (conv1): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm1): GroupNorm(2, 64, eps=1e-05, affine=True)
          (conv2): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm2): GroupNorm(2, 64, eps=1e-05, affine=True)
          (shortcut): Identity()
        )
      )
      (downsample): Conv3d(64, 64, kernel_size=(2, 2, 2), stride=(2, 2, 2))
    )
    (1): DownBlock3D(
      (resnets): ModuleList(
        (0): ResnetBlock3D(
          (conv1): Conv3d(64, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm1): GroupNorm(4, 128, eps=1e-05, affine=True)
          (conv2): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm2): GroupNorm(4, 128, eps=1e-05, affine=True)
          (shortcut): Conv3d(64, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))
        )
        (1): ResnetBlock3D(
          (conv1): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm1): GroupNorm(4, 128, eps=1e-05, affine=True)
          (conv2): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm2): GroupNorm(4, 128, eps=1e-05, affine=True)
          (shortcut): Identity()
        )
      )
      (downsample): Conv3d(128, 128, kernel_size=(2, 2, 2), stride=(2, 2, 2))
    )
    (2): DownBlock3D(
      (resnets): ModuleList(
        (0): ResnetBlock3D(
          (conv1): Conv3d(128, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm1): GroupNorm(8, 256, eps=1e-05, affine=True)
          (conv2): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm2): GroupNorm(8, 256, eps=1e-05, affine=True)
          (shortcut): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 1, 1))
        )
        (1): ResnetBlock3D(
          (conv1): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm1): GroupNorm(8, 256, eps=1e-05, affine=True)
          (conv2): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm2): GroupNorm(8, 256, eps=1e-05, affine=True)
          (shortcut): Identity()
        )
      )
      (downsample): Conv3d(256, 256, kernel_size=(2, 2, 2), stride=(2, 2, 2))
    )
  )
  (mid_resnet1): ResnetBlock3D(
    (conv1): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
    (norm1): GroupNorm(8, 256, eps=1e-05, affine=True)
    (conv2): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
    (norm2): GroupNorm(8, 256, eps=1e-05, affine=True)
    (shortcut): Identity()
  )
  (mid_resnet2): ResnetBlock3D(
    (conv1): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
    (norm1): GroupNorm(8, 256, eps=1e-05, affine=True)
    (conv2): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
    (norm2): GroupNorm(8, 256, eps=1e-05, affine=True)
    (shortcut): Identity()
  )
  (up_blocks): ModuleList(
    (0): UpBlock3D(
      (resnets): ModuleList(
        (0): ResnetBlock3D(
          (conv1): Conv3d(512, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm1): GroupNorm(4, 128, eps=1e-05, affine=True)
          (conv2): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm2): GroupNorm(4, 128, eps=1e-05, affine=True)
          (shortcut): Conv3d(512, 128, kernel_size=(1, 1, 1), stride=(1, 1, 1))
        )
        (1): ResnetBlock3D(
          (conv1): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm1): GroupNorm(4, 128, eps=1e-05, affine=True)
          (conv2): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm2): GroupNorm(4, 128, eps=1e-05, affine=True)
          (shortcut): Identity()
        )
      )
      (upsample): ConvTranspose3d(256, 256, kernel_size=(2, 2, 2), stride=(2, 2, 2))
    )
    (1): UpBlock3D(
      (resnets): ModuleList(
        (0): ResnetBlock3D(
          (conv1): Conv3d(256, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm1): GroupNorm(2, 64, eps=1e-05, affine=True)
          (conv2): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm2): GroupNorm(2, 64, eps=1e-05, affine=True)
          (shortcut): Conv3d(256, 64, kernel_size=(1, 1, 1), stride=(1, 1, 1))
        )
        (1): ResnetBlock3D(
          (conv1): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm1): GroupNorm(2, 64, eps=1e-05, affine=True)
          (conv2): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm2): GroupNorm(2, 64, eps=1e-05, affine=True)
          (shortcut): Identity()
        )
      )
      (upsample): ConvTranspose3d(128, 128, kernel_size=(2, 2, 2), stride=(2, 2, 2))
    )
    (2): UpBlock3D(
      (resnets): ModuleList(
        (0): ResnetBlock3D(
          (conv1): Conv3d(128, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm1): GroupNorm(1, 32, eps=1e-05, affine=True)
          (conv2): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm2): GroupNorm(1, 32, eps=1e-05, affine=True)
          (shortcut): Conv3d(128, 32, kernel_size=(1, 1, 1), stride=(1, 1, 1))
        )
        (1): ResnetBlock3D(
          (conv1): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm1): GroupNorm(1, 32, eps=1e-05, affine=True)
          (conv2): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
          (norm2): GroupNorm(1, 32, eps=1e-05, affine=True)
          (shortcut): Identity()
        )
      )
      (upsample): ConvTranspose3d(64, 64, kernel_size=(2, 2, 2), stride=(2, 2, 2))
    )
  )
  (conv_out): Conv3d(32, 40, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
)

UNet2D

PaddedUNet2D(
  (conv_in): Conv2d(84, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (down_blocks): ModuleList(
    (0): DownBlock2D(
      (resnets): ModuleList(
        (0): ResnetBlock2D(
          (conv1): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm1): GroupNorm(2, 64, eps=1e-05, affine=True)
          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm2): GroupNorm(2, 64, eps=1e-05, affine=True)
          (shortcut): Conv2d(32, 64, kernel_size=(1, 1), stride=(1, 1))
        )
        (1): ResnetBlock2D(
          (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm1): GroupNorm(2, 64, eps=1e-05, affine=True)
          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm2): GroupNorm(2, 64, eps=1e-05, affine=True)
          (shortcut): Identity()
        )
      )
      (downsample): Conv2d(64, 64, kernel_size=(2, 2), stride=(2, 2))
    )
    (1): DownBlock2D(
      (resnets): ModuleList(
        (0): ResnetBlock2D(
          (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm1): GroupNorm(4, 128, eps=1e-05, affine=True)
          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm2): GroupNorm(4, 128, eps=1e-05, affine=True)
          (shortcut): Conv2d(64, 128, kernel_size=(1, 1), stride=(1, 1))
        )
        (1): ResnetBlock2D(
          (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm1): GroupNorm(4, 128, eps=1e-05, affine=True)
          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm2): GroupNorm(4, 128, eps=1e-05, affine=True)
          (shortcut): Identity()
        )
      )
      (downsample): Conv2d(128, 128, kernel_size=(2, 2), stride=(2, 2))
    )
    (2): DownBlock2D(
      (resnets): ModuleList(
        (0): ResnetBlock2D(
          (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm1): GroupNorm(8, 256, eps=1e-05, affine=True)
          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm2): GroupNorm(8, 256, eps=1e-05, affine=True)
          (shortcut): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1))
        )
        (1): ResnetBlock2D(
          (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm1): GroupNorm(8, 256, eps=1e-05, affine=True)
          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm2): GroupNorm(8, 256, eps=1e-05, affine=True)
          (shortcut): Identity()
        )
      )
      (downsample): Conv2d(256, 256, kernel_size=(2, 2), stride=(2, 2))
    )
  )
  (mid_resnet1): ResnetBlock2D(
    (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (norm1): GroupNorm(8, 256, eps=1e-05, affine=True)
    (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (norm2): GroupNorm(8, 256, eps=1e-05, affine=True)
    (shortcut): Identity()
  )
  (mid_resnet2): ResnetBlock2D(
    (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (norm1): GroupNorm(8, 256, eps=1e-05, affine=True)
    (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (norm2): GroupNorm(8, 256, eps=1e-05, affine=True)
    (shortcut): Identity()
  )
  (up_blocks): ModuleList(
    (0): UpBlock2D(
      (resnets): ModuleList(
        (0): ResnetBlock2D(
          (conv1): Conv2d(512, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm1): GroupNorm(4, 128, eps=1e-05, affine=True)
          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm2): GroupNorm(4, 128, eps=1e-05, affine=True)
          (shortcut): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1))
        )
        (1): ResnetBlock2D(
          (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm1): GroupNorm(4, 128, eps=1e-05, affine=True)
          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm2): GroupNorm(4, 128, eps=1e-05, affine=True)
          (shortcut): Identity()
        )
      )
      (upsample): ConvTranspose2d(256, 256, kernel_size=(2, 2), stride=(2, 2))
    )
    (1): UpBlock2D(
      (resnets): ModuleList(
        (0): ResnetBlock2D(
          (conv1): Conv2d(256, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm1): GroupNorm(2, 64, eps=1e-05, affine=True)
          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm2): GroupNorm(2, 64, eps=1e-05, affine=True)
          (shortcut): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1))
        )
        (1): ResnetBlock2D(
          (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm1): GroupNorm(2, 64, eps=1e-05, affine=True)
          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm2): GroupNorm(2, 64, eps=1e-05, affine=True)
          (shortcut): Identity()
        )
      )
      (upsample): ConvTranspose2d(128, 128, kernel_size=(2, 2), stride=(2, 2))
    )
    (2): UpBlock2D(
      (resnets): ModuleList(
        (0): ResnetBlock2D(
          (conv1): Conv2d(128, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm1): GroupNorm(1, 32, eps=1e-05, affine=True)
          (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm2): GroupNorm(1, 32, eps=1e-05, affine=True)
          (shortcut): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1))
        )
        (1): ResnetBlock2D(
          (conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm1): GroupNorm(1, 32, eps=1e-05, affine=True)
          (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (norm2): GroupNorm(1, 32, eps=1e-05, affine=True)
          (shortcut): Identity()
        )
      )
      (upsample): ConvTranspose2d(64, 64, kernel_size=(2, 2), stride=(2, 2))
    )
  )
  (conv_out): Conv2d(32, 80, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
)