Skip to content

Instantly share code, notes, and snippets.

@laksjdjf

laksjdjf/UNet_layer.txt

Created April 3, 2023 14:42
Show Gist options
  • Save laksjdjf/d3eba3e58351188bea84c90eebb122c9 to your computer and use it in GitHub Desktop.
Save laksjdjf/d3eba3e58351188bea84c90eebb122c9 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
UNet_layer.txt
==============================================================================================================
Layer (type (var_name)) Input Shape Output Shape
==============================================================================================================
UNetModel (UNetModel) [1, 4, 64, 64] [1, 4, 64, 64]
├─Sequential (time_embed) [1, 320] [1, 1280]
│ └─Linear (0) [1, 320] [1, 1280]
│ └─SiLU (1) [1, 1280] [1, 1280]
│ └─Linear (2) [1, 1280] [1, 1280]
├─ModuleList (input_blocks) -- --
│ └─TimestepEmbedSequential (0) [1, 4, 64, 64] [1, 320, 64, 64]
│ │ └─Conv2d (0) [1, 4, 64, 64] [1, 320, 64, 64]
│ └─TimestepEmbedSequential (1) [1, 320, 64, 64] [1, 320, 64, 64]
│ │ └─ResBlock (0) [1, 320, 64, 64] [1, 320, 64, 64]
│ │ └─SpatialTransformer (1) [1, 320, 64, 64] [1, 320, 64, 64]
│ └─TimestepEmbedSequential (2) [1, 320, 64, 64] [1, 320, 64, 64]
│ │ └─ResBlock (0) [1, 320, 64, 64] [1, 320, 64, 64]
│ │ └─SpatialTransformer (1) [1, 320, 64, 64] [1, 320, 64, 64]
│ └─TimestepEmbedSequential (3) [1, 320, 64, 64] [1, 320, 32, 32]
│ │ └─Downsample (0) [1, 320, 64, 64] [1, 320, 32, 32]
│ └─TimestepEmbedSequential (4) [1, 320, 32, 32] [1, 640, 32, 32]
│ │ └─ResBlock (0) [1, 320, 32, 32] [1, 640, 32, 32]
│ │ └─SpatialTransformer (1) [1, 640, 32, 32] [1, 640, 32, 32]
│ └─TimestepEmbedSequential (5) [1, 640, 32, 32] [1, 640, 32, 32]
│ │ └─ResBlock (0) [1, 640, 32, 32] [1, 640, 32, 32]
│ │ └─SpatialTransformer (1) [1, 640, 32, 32] [1, 640, 32, 32]
│ └─TimestepEmbedSequential (6) [1, 640, 32, 32] [1, 640, 16, 16]
│ │ └─Downsample (0) [1, 640, 32, 32] [1, 640, 16, 16]
│ └─TimestepEmbedSequential (7) [1, 640, 16, 16] [1, 1280, 16, 16]
│ │ └─ResBlock (0) [1, 640, 16, 16] [1, 1280, 16, 16]
│ │ └─SpatialTransformer (1) [1, 1280, 16, 16] [1, 1280, 16, 16]
│ └─TimestepEmbedSequential (8) [1, 1280, 16, 16] [1, 1280, 16, 16]
│ │ └─ResBlock (0) [1, 1280, 16, 16] [1, 1280, 16, 16]
│ │ └─SpatialTransformer (1) [1, 1280, 16, 16] [1, 1280, 16, 16]
│ └─TimestepEmbedSequential (9) [1, 1280, 16, 16] [1, 1280, 8, 8]
│ │ └─Downsample (0) [1, 1280, 16, 16] [1, 1280, 8, 8]
│ └─TimestepEmbedSequential (10) [1, 1280, 8, 8] [1, 1280, 8, 8]
│ │ └─ResBlock (0) [1, 1280, 8, 8] [1, 1280, 8, 8]
│ └─TimestepEmbedSequential (11) [1, 1280, 8, 8] [1, 1280, 8, 8]
│ │ └─ResBlock (0) [1, 1280, 8, 8] [1, 1280, 8, 8]
├─TimestepEmbedSequential (middle_block) [1, 1280, 8, 8] [1, 1280, 8, 8]
│ └─ResBlock (0) [1, 1280, 8, 8] [1, 1280, 8, 8]
│ │ └─Sequential (in_layers) [1, 1280, 8, 8] [1, 1280, 8, 8]
│ │ └─Sequential (emb_layers) [1, 1280] [1, 1280]
│ │ └─Sequential (out_layers) [1, 1280, 8, 8] [1, 1280, 8, 8]
│ │ └─Identity (skip_connection) [1, 1280, 8, 8] [1, 1280, 8, 8]
│ └─SpatialTransformer (1) [1, 1280, 8, 8] [1, 1280, 8, 8]
│ │ └─GroupNorm (norm) [1, 1280, 8, 8] [1, 1280, 8, 8]
│ │ └─Linear (proj_in) [1, 64, 1280] [1, 64, 1280]
│ │ └─ModuleList (transformer_blocks) -- --
│ │ └─Linear (proj_out) [1, 64, 1280] [1, 64, 1280]
│ └─ResBlock (2) [1, 1280, 8, 8] [1, 1280, 8, 8]
│ │ └─Sequential (in_layers) [1, 1280, 8, 8] [1, 1280, 8, 8]
│ │ └─Sequential (emb_layers) [1, 1280] [1, 1280]
│ │ └─Sequential (out_layers) [1, 1280, 8, 8] [1, 1280, 8, 8]
│ │ └─Identity (skip_connection) [1, 1280, 8, 8] [1, 1280, 8, 8]
├─ModuleList (output_blocks) -- --
│ └─TimestepEmbedSequential (0) [1, 2560, 8, 8] [1, 1280, 8, 8]
│ │ └─ResBlock (0) [1, 2560, 8, 8] [1, 1280, 8, 8]
│ └─TimestepEmbedSequential (1) [1, 2560, 8, 8] [1, 1280, 8, 8]
│ │ └─ResBlock (0) [1, 2560, 8, 8] [1, 1280, 8, 8]
│ └─TimestepEmbedSequential (2) [1, 2560, 8, 8] [1, 1280, 16, 16]
│ │ └─ResBlock (0) [1, 2560, 8, 8] [1, 1280, 8, 8]
│ │ └─Upsample (1) [1, 1280, 8, 8] [1, 1280, 16, 16]
│ └─TimestepEmbedSequential (3) [1, 2560, 16, 16] [1, 1280, 16, 16]
│ │ └─ResBlock (0) [1, 2560, 16, 16] [1, 1280, 16, 16]
│ │ └─SpatialTransformer (1) [1, 1280, 16, 16] [1, 1280, 16, 16]
│ └─TimestepEmbedSequential (4) [1, 2560, 16, 16] [1, 1280, 16, 16]
│ │ └─ResBlock (0) [1, 2560, 16, 16] [1, 1280, 16, 16]
│ │ └─SpatialTransformer (1) [1, 1280, 16, 16] [1, 1280, 16, 16]
│ └─TimestepEmbedSequential (5) [1, 1920, 16, 16] [1, 1280, 32, 32]
│ │ └─ResBlock (0) [1, 1920, 16, 16] [1, 1280, 16, 16]
│ │ └─SpatialTransformer (1) [1, 1280, 16, 16] [1, 1280, 16, 16]
│ │ └─Upsample (2) [1, 1280, 16, 16] [1, 1280, 32, 32]
│ └─TimestepEmbedSequential (6) [1, 1920, 32, 32] [1, 640, 32, 32]
│ │ └─ResBlock (0) [1, 1920, 32, 32] [1, 640, 32, 32]
│ │ └─SpatialTransformer (1) [1, 640, 32, 32] [1, 640, 32, 32]
│ └─TimestepEmbedSequential (7) [1, 1280, 32, 32] [1, 640, 32, 32]
│ │ └─ResBlock (0) [1, 1280, 32, 32] [1, 640, 32, 32]
│ │ └─SpatialTransformer (1) [1, 640, 32, 32] [1, 640, 32, 32]
│ └─TimestepEmbedSequential (8) [1, 960, 32, 32] [1, 640, 64, 64]
│ │ └─ResBlock (0) [1, 960, 32, 32] [1, 640, 32, 32]
│ │ └─SpatialTransformer (1) [1, 640, 32, 32] [1, 640, 32, 32]
│ │ └─Upsample (2) [1, 640, 32, 32] [1, 640, 64, 64]
│ └─TimestepEmbedSequential (9) [1, 960, 64, 64] [1, 320, 64, 64]
│ │ └─ResBlock (0) [1, 960, 64, 64] [1, 320, 64, 64]
│ │ └─SpatialTransformer (1) [1, 320, 64, 64] [1, 320, 64, 64]
│ └─TimestepEmbedSequential (10) [1, 640, 64, 64] [1, 320, 64, 64]
│ │ └─ResBlock (0) [1, 640, 64, 64] [1, 320, 64, 64]
│ │ └─SpatialTransformer (1) [1, 320, 64, 64] [1, 320, 64, 64]
│ └─TimestepEmbedSequential (11) [1, 640, 64, 64] [1, 320, 64, 64]
│ │ └─ResBlock (0) [1, 640, 64, 64] [1, 320, 64, 64]
│ │ └─SpatialTransformer (1) [1, 320, 64, 64] [1, 320, 64, 64]
├─Sequential (out) [1, 320, 64, 64] [1, 4, 64, 64]
│ └─GroupNorm32 (0) [1, 320, 64, 64] [1, 320, 64, 64]
│ └─SiLU (1) [1, 320, 64, 64] [1, 320, 64, 64]
│ └─Conv2d (2) [1, 320, 64, 64] [1, 4, 64, 64]
==============================================================================================================
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment