tutorials/models_cls_seg.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Models Classification Segmentation from libraires    \n",
    "* Platform: Windows, Python, miniconda, PyTorch    \n",
    "* Date: 2024-09-20    \n",
    "* Author: Jing Zhang    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## models from TIMM\n",
    "* Reference: https://huggingface.co/timm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "list models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bat_resnext26ts\n",
      "beit_base_patch16_224\n",
      "beit_base_patch16_384\n",
      "beit_large_patch16_224\n",
      "beit_large_patch16_384\n",
      "beit_large_patch16_512\n",
      "beitv2_base_patch16_224\n",
      "beitv2_large_patch16_224\n",
      "botnet26t_256\n",
      "botnet50ts_256\n",
      "caformer_b36\n",
      "caformer_m36\n",
      "caformer_s18\n",
      "caformer_s36\n",
      "cait_m36_384\n",
      "cait_m48_448\n",
      "cait_s24_224\n",
      "cait_s24_384\n",
      "cait_s36_384\n",
      "cait_xs24_384\n",
      "cait_xxs24_224\n",
      "cait_xxs24_384\n",
      "cait_xxs36_224\n",
      "cait_xxs36_384\n",
      "coat_lite_medium\n",
      "coat_lite_medium_384\n",
      "coat_lite_mini\n",
      "coat_lite_small\n",
      "coat_lite_tiny\n",
      "coat_mini\n",
      "coat_small\n",
      "coat_tiny\n",
      "coatnet_0_224\n",
      "coatnet_0_rw_224\n",
      "coatnet_1_224\n",
      "coatnet_1_rw_224\n",
      "coatnet_2_224\n",
      "coatnet_2_rw_224\n",
      "coatnet_3_224\n",
      "coatnet_3_rw_224\n",
      "coatnet_4_224\n",
      "coatnet_5_224\n",
      "coatnet_bn_0_rw_224\n",
      "coatnet_nano_cc_224\n",
      "coatnet_nano_rw_224\n",
      "coatnet_pico_rw_224\n",
      "coatnet_rmlp_0_rw_224\n",
      "coatnet_rmlp_1_rw2_224\n",
      "coatnet_rmlp_1_rw_224\n",
      "coatnet_rmlp_2_rw_224\n",
      "coatnet_rmlp_2_rw_384\n",
      "coatnet_rmlp_3_rw_224\n",
      "coatnet_rmlp_nano_rw_224\n",
      "coatnext_nano_rw_224\n",
      "convformer_b36\n",
      "convformer_m36\n",
      "convformer_s18\n",
      "convformer_s36\n",
      "convit_base\n",
      "convit_small\n",
      "convit_tiny\n",
      "convmixer_768_32\n",
      "convmixer_1024_20_ks9_p14\n",
      "convmixer_1536_20\n",
      "convnext_atto\n",
      "convnext_atto_ols\n",
      "convnext_base\n",
      "convnext_femto\n",
      "convnext_femto_ols\n",
      "convnext_large\n",
      "convnext_large_mlp\n",
      "convnext_nano\n",
      "convnext_nano_ols\n",
      "convnext_pico\n",
      "convnext_pico_ols\n",
      "convnext_small\n",
      "convnext_tiny\n",
      "convnext_tiny_hnf\n",
      "convnext_xlarge\n",
      "convnext_xxlarge\n",
      "convnextv2_atto\n",
      "convnextv2_base\n",
      "convnextv2_femto\n",
      "convnextv2_huge\n",
      "convnextv2_large\n",
      "convnextv2_nano\n",
      "convnextv2_pico\n",
      "convnextv2_small\n",
      "convnextv2_tiny\n",
      "crossvit_9_240\n",
      "crossvit_9_dagger_240\n",
      "crossvit_15_240\n",
      "crossvit_15_dagger_240\n",
      "crossvit_15_dagger_408\n",
      "crossvit_18_240\n",
      "crossvit_18_dagger_240\n",
      "crossvit_18_dagger_408\n",
      "crossvit_base_240\n",
      "crossvit_small_240\n",
      "crossvit_tiny_240\n",
      "cs3darknet_focus_l\n",
      "cs3darknet_focus_m\n",
      "cs3darknet_focus_s\n",
      "cs3darknet_focus_x\n",
      "cs3darknet_l\n",
      "cs3darknet_m\n",
      "cs3darknet_s\n",
      "cs3darknet_x\n",
      "cs3edgenet_x\n",
      "cs3se_edgenet_x\n",
      "cs3sedarknet_l\n",
      "cs3sedarknet_x\n",
      "cs3sedarknet_xdw\n",
      "cspdarknet53\n",
      "cspresnet50\n",
      "cspresnet50d\n",
      "cspresnet50w\n",
      "cspresnext50\n",
      "darknet17\n",
      "darknet21\n",
      "darknet53\n",
      "darknetaa53\n",
      "davit_base\n",
      "davit_giant\n",
      "davit_huge\n",
      "davit_large\n",
      "davit_small\n",
      "davit_tiny\n",
      "deit3_base_patch16_224\n",
      "deit3_base_patch16_384\n",
      "deit3_huge_patch14_224\n",
      "deit3_large_patch16_224\n",
      "deit3_large_patch16_384\n",
      "deit3_medium_patch16_224\n",
      "deit3_small_patch16_224\n",
      "deit3_small_patch16_384\n",
      "deit_base_distilled_patch16_224\n",
      "deit_base_distilled_patch16_384\n",
      "deit_base_patch16_224\n",
      "deit_base_patch16_384\n",
      "deit_small_distilled_patch16_224\n",
      "deit_small_patch16_224\n",
      "deit_tiny_distilled_patch16_224\n",
      "deit_tiny_patch16_224\n",
      "densenet121\n",
      "densenet161\n",
      "densenet169\n",
      "densenet201\n",
      "densenet264d\n",
      "densenetblur121d\n",
      "dla34\n",
      "dla46_c\n",
      "dla46x_c\n",
      "dla60\n",
      "dla60_res2net\n",
      "dla60_res2next\n",
      "dla60x\n",
      "dla60x_c\n",
      "dla102\n",
      "dla102x\n",
      "dla102x2\n",
      "dla169\n",
      "dm_nfnet_f0\n",
      "dm_nfnet_f1\n",
      "dm_nfnet_f2\n",
      "dm_nfnet_f3\n",
      "dm_nfnet_f4\n",
      "dm_nfnet_f5\n",
      "dm_nfnet_f6\n",
      "dpn48b\n",
      "dpn68\n",
      "dpn68b\n",
      "dpn92\n",
      "dpn98\n",
      "dpn107\n",
      "dpn131\n",
      "eca_botnext26ts_256\n",
      "eca_halonext26ts\n",
      "eca_nfnet_l0\n",
      "eca_nfnet_l1\n",
      "eca_nfnet_l2\n",
      "eca_nfnet_l3\n",
      "eca_resnet33ts\n",
      "eca_resnext26ts\n",
      "eca_vovnet39b\n",
      "ecaresnet26t\n",
      "ecaresnet50d\n",
      "ecaresnet50d_pruned\n",
      "ecaresnet50t\n",
      "ecaresnet101d\n",
      "ecaresnet101d_pruned\n",
      "ecaresnet200d\n",
      "ecaresnet269d\n",
      "ecaresnetlight\n",
      "ecaresnext26t_32x4d\n",
      "ecaresnext50t_32x4d\n",
      "edgenext_base\n",
      "edgenext_small\n",
      "edgenext_small_rw\n",
      "edgenext_x_small\n",
      "edgenext_xx_small\n",
      "efficientformer_l1\n",
      "efficientformer_l3\n",
      "efficientformer_l7\n",
      "efficientformerv2_l\n",
      "efficientformerv2_s0\n",
      "efficientformerv2_s1\n",
      "efficientformerv2_s2\n",
      "efficientnet_b0\n",
      "efficientnet_b0_g8_gn\n",
      "efficientnet_b0_g16_evos\n",
      "efficientnet_b0_gn\n",
      "efficientnet_b1\n",
      "efficientnet_b1_pruned\n",
      "efficientnet_b2\n",
      "efficientnet_b2_pruned\n",
      "efficientnet_b3\n",
      "efficientnet_b3_g8_gn\n",
      "efficientnet_b3_gn\n",
      "efficientnet_b3_pruned\n",
      "efficientnet_b4\n",
      "efficientnet_b5\n",
      "efficientnet_b6\n",
      "efficientnet_b7\n",
      "efficientnet_b8\n",
      "efficientnet_cc_b0_4e\n",
      "efficientnet_cc_b0_8e\n",
      "efficientnet_cc_b1_8e\n",
      "efficientnet_el\n",
      "efficientnet_el_pruned\n",
      "efficientnet_em\n",
      "efficientnet_es\n",
      "efficientnet_es_pruned\n",
      "efficientnet_l2\n",
      "efficientnet_lite0\n",
      "efficientnet_lite1\n",
      "efficientnet_lite2\n",
      "efficientnet_lite3\n",
      "efficientnet_lite4\n",
      "efficientnetv2_l\n",
      "efficientnetv2_m\n",
      "efficientnetv2_rw_m\n",
      "efficientnetv2_rw_s\n",
      "efficientnetv2_rw_t\n",
      "efficientnetv2_s\n",
      "efficientnetv2_xl\n",
      "efficientvit_b0\n",
      "efficientvit_b1\n",
      "efficientvit_b2\n",
      "efficientvit_b3\n",
      "efficientvit_l1\n",
      "efficientvit_l2\n",
      "efficientvit_l3\n",
      "efficientvit_m0\n",
      "efficientvit_m1\n",
      "efficientvit_m2\n",
      "efficientvit_m3\n",
      "efficientvit_m4\n",
      "efficientvit_m5\n",
      "ese_vovnet19b_dw\n",
      "ese_vovnet19b_slim\n",
      "ese_vovnet19b_slim_dw\n",
      "ese_vovnet39b\n",
      "ese_vovnet39b_evos\n",
      "ese_vovnet57b\n",
      "ese_vovnet99b\n",
      "eva02_base_patch14_224\n",
      "eva02_base_patch14_448\n",
      "eva02_base_patch16_clip_224\n",
      "eva02_enormous_patch14_clip_224\n",
      "eva02_large_patch14_224\n",
      "eva02_large_patch14_448\n",
      "eva02_large_patch14_clip_224\n",
      "eva02_large_patch14_clip_336\n",
      "eva02_small_patch14_224\n",
      "eva02_small_patch14_336\n",
      "eva02_tiny_patch14_224\n",
      "eva02_tiny_patch14_336\n",
      "eva_giant_patch14_224\n",
      "eva_giant_patch14_336\n",
      "eva_giant_patch14_560\n",
      "eva_giant_patch14_clip_224\n",
      "eva_large_patch14_196\n",
      "eva_large_patch14_336\n",
      "fastvit_ma36\n",
      "fastvit_s12\n",
      "fastvit_sa12\n",
      "fastvit_sa24\n",
      "fastvit_sa36\n",
      "fastvit_t8\n",
      "fastvit_t12\n",
      "fbnetc_100\n",
      "fbnetv3_b\n",
      "fbnetv3_d\n",
      "fbnetv3_g\n",
      "flexivit_base\n",
      "flexivit_large\n",
      "flexivit_small\n",
      "focalnet_base_lrf\n",
      "focalnet_base_srf\n",
      "focalnet_huge_fl3\n",
      "focalnet_huge_fl4\n",
      "focalnet_large_fl3\n",
      "focalnet_large_fl4\n",
      "focalnet_small_lrf\n",
      "focalnet_small_srf\n",
      "focalnet_tiny_lrf\n",
      "focalnet_tiny_srf\n",
      "focalnet_xlarge_fl3\n",
      "focalnet_xlarge_fl4\n",
      "gc_efficientnetv2_rw_t\n",
      "gcresnet33ts\n",
      "gcresnet50t\n",
      "gcresnext26ts\n",
      "gcresnext50ts\n",
      "gcvit_base\n",
      "gcvit_small\n",
      "gcvit_tiny\n",
      "gcvit_xtiny\n",
      "gcvit_xxtiny\n",
      "gernet_l\n",
      "gernet_m\n",
      "gernet_s\n",
      "ghostnet_050\n",
      "ghostnet_100\n",
      "ghostnet_130\n",
      "ghostnetv2_100\n",
      "ghostnetv2_130\n",
      "ghostnetv2_160\n",
      "gmixer_12_224\n",
      "gmixer_24_224\n",
      "gmlp_b16_224\n",
      "gmlp_s16_224\n",
      "gmlp_ti16_224\n",
      "halo2botnet50ts_256\n",
      "halonet26t\n",
      "halonet50ts\n",
      "halonet_h1\n",
      "haloregnetz_b\n",
      "hardcorenas_a\n",
      "hardcorenas_b\n",
      "hardcorenas_c\n",
      "hardcorenas_d\n",
      "hardcorenas_e\n",
      "hardcorenas_f\n",
      "hgnet_base\n",
      "hgnet_small\n",
      "hgnet_tiny\n",
      "hgnetv2_b0\n",
      "hgnetv2_b1\n",
      "hgnetv2_b2\n",
      "hgnetv2_b3\n",
      "hgnetv2_b4\n",
      "hgnetv2_b5\n",
      "hgnetv2_b6\n",
      "hrnet_w18\n",
      "hrnet_w18_small\n",
      "hrnet_w18_small_v2\n",
      "hrnet_w18_ssld\n",
      "hrnet_w30\n",
      "hrnet_w32\n",
      "hrnet_w40\n",
      "hrnet_w44\n",
      "hrnet_w48\n",
      "hrnet_w48_ssld\n",
      "hrnet_w64\n",
      "inception_next_base\n",
      "inception_next_small\n",
      "inception_next_tiny\n",
      "inception_resnet_v2\n",
      "inception_v3\n",
      "inception_v4\n",
      "lambda_resnet26rpt_256\n",
      "lambda_resnet26t\n",
      "lambda_resnet50ts\n",
      "lamhalobotnet50ts_256\n",
      "lcnet_035\n",
      "lcnet_050\n",
      "lcnet_075\n",
      "lcnet_100\n",
      "lcnet_150\n",
      "legacy_senet154\n",
      "legacy_seresnet18\n",
      "legacy_seresnet34\n",
      "legacy_seresnet50\n",
      "legacy_seresnet101\n",
      "legacy_seresnet152\n",
      "legacy_seresnext26_32x4d\n",
      "legacy_seresnext50_32x4d\n",
      "legacy_seresnext101_32x4d\n",
      "legacy_xception\n",
      "levit_128\n",
      "levit_128s\n",
      "levit_192\n",
      "levit_256\n",
      "levit_256d\n",
      "levit_384\n",
      "levit_384_s8\n",
      "levit_512\n",
      "levit_512_s8\n",
      "levit_512d\n",
      "levit_conv_128\n",
      "levit_conv_128s\n",
      "levit_conv_192\n",
      "levit_conv_256\n",
      "levit_conv_256d\n",
      "levit_conv_384\n",
      "levit_conv_384_s8\n",
      "levit_conv_512\n",
      "levit_conv_512_s8\n",
      "levit_conv_512d\n",
      "maxvit_base_tf_224\n",
      "maxvit_base_tf_384\n",
      "maxvit_base_tf_512\n",
      "maxvit_large_tf_224\n",
      "maxvit_large_tf_384\n",
      "maxvit_large_tf_512\n",
      "maxvit_nano_rw_256\n",
      "maxvit_pico_rw_256\n",
      "maxvit_rmlp_base_rw_224\n",
      "maxvit_rmlp_base_rw_384\n",
      "maxvit_rmlp_nano_rw_256\n",
      "maxvit_rmlp_pico_rw_256\n",
      "maxvit_rmlp_small_rw_224\n",
      "maxvit_rmlp_small_rw_256\n",
      "maxvit_rmlp_tiny_rw_256\n",
      "maxvit_small_tf_224\n",
      "maxvit_small_tf_384\n",
      "maxvit_small_tf_512\n",
      "maxvit_tiny_pm_256\n",
      "maxvit_tiny_rw_224\n",
      "maxvit_tiny_rw_256\n",
      "maxvit_tiny_tf_224\n",
      "maxvit_tiny_tf_384\n",
      "maxvit_tiny_tf_512\n",
      "maxvit_xlarge_tf_224\n",
      "maxvit_xlarge_tf_384\n",
      "maxvit_xlarge_tf_512\n",
      "maxxvit_rmlp_nano_rw_256\n",
      "maxxvit_rmlp_small_rw_256\n",
      "maxxvit_rmlp_tiny_rw_256\n",
      "maxxvitv2_nano_rw_256\n",
      "maxxvitv2_rmlp_base_rw_224\n",
      "maxxvitv2_rmlp_base_rw_384\n",
      "maxxvitv2_rmlp_large_rw_224\n",
      "mixer_b16_224\n",
      "mixer_b32_224\n",
      "mixer_l16_224\n",
      "mixer_l32_224\n",
      "mixer_s16_224\n",
      "mixer_s32_224\n",
      "mixnet_l\n",
      "mixnet_m\n",
      "mixnet_s\n",
      "mixnet_xl\n",
      "mixnet_xxl\n",
      "mnasnet_050\n",
      "mnasnet_075\n",
      "mnasnet_100\n",
      "mnasnet_140\n",
      "mnasnet_small\n",
      "mobilenetv2_035\n",
      "mobilenetv2_050\n",
      "mobilenetv2_075\n",
      "mobilenetv2_100\n",
      "mobilenetv2_110d\n",
      "mobilenetv2_120d\n",
      "mobilenetv2_140\n",
      "mobilenetv3_large_075\n",
      "mobilenetv3_large_100\n",
      "mobilenetv3_rw\n",
      "mobilenetv3_small_050\n",
      "mobilenetv3_small_075\n",
      "mobilenetv3_small_100\n",
      "mobileone_s0\n",
      "mobileone_s1\n",
      "mobileone_s2\n",
      "mobileone_s3\n",
      "mobileone_s4\n",
      "mobilevit_s\n",
      "mobilevit_xs\n",
      "mobilevit_xxs\n",
      "mobilevitv2_050\n",
      "mobilevitv2_075\n",
      "mobilevitv2_100\n",
      "mobilevitv2_125\n",
      "mobilevitv2_150\n",
      "mobilevitv2_175\n",
      "mobilevitv2_200\n",
      "mvitv2_base\n",
      "mvitv2_base_cls\n",
      "mvitv2_huge_cls\n",
      "mvitv2_large\n",
      "mvitv2_large_cls\n",
      "mvitv2_small\n",
      "mvitv2_small_cls\n",
      "mvitv2_tiny\n",
      "nasnetalarge\n",
      "nest_base\n",
      "nest_base_jx\n",
      "nest_small\n",
      "nest_small_jx\n",
      "nest_tiny\n",
      "nest_tiny_jx\n",
      "nextvit_base\n",
      "nextvit_large\n",
      "nextvit_small\n",
      "nf_ecaresnet26\n",
      "nf_ecaresnet50\n",
      "nf_ecaresnet101\n",
      "nf_regnet_b0\n",
      "nf_regnet_b1\n",
      "nf_regnet_b2\n",
      "nf_regnet_b3\n",
      "nf_regnet_b4\n",
      "nf_regnet_b5\n",
      "nf_resnet26\n",
      "nf_resnet50\n",
      "nf_resnet101\n",
      "nf_seresnet26\n",
      "nf_seresnet50\n",
      "nf_seresnet101\n",
      "nfnet_f0\n",
      "nfnet_f1\n",
      "nfnet_f2\n",
      "nfnet_f3\n",
      "nfnet_f4\n",
      "nfnet_f5\n",
      "nfnet_f6\n",
      "nfnet_f7\n",
      "nfnet_l0\n",
      "pit_b_224\n",
      "pit_b_distilled_224\n",
      "pit_s_224\n",
      "pit_s_distilled_224\n",
      "pit_ti_224\n",
      "pit_ti_distilled_224\n",
      "pit_xs_224\n",
      "pit_xs_distilled_224\n",
      "pnasnet5large\n",
      "poolformer_m36\n",
      "poolformer_m48\n",
      "poolformer_s12\n",
      "poolformer_s24\n",
      "poolformer_s36\n",
      "poolformerv2_m36\n",
      "poolformerv2_m48\n",
      "poolformerv2_s12\n",
      "poolformerv2_s24\n",
      "poolformerv2_s36\n",
      "pvt_v2_b0\n",
      "pvt_v2_b1\n",
      "pvt_v2_b2\n",
      "pvt_v2_b2_li\n",
      "pvt_v2_b3\n",
      "pvt_v2_b4\n",
      "pvt_v2_b5\n",
      "regnetv_040\n",
      "regnetv_064\n",
      "regnetx_002\n",
      "regnetx_004\n",
      "regnetx_004_tv\n",
      "regnetx_006\n",
      "regnetx_008\n",
      "regnetx_016\n",
      "regnetx_032\n",
      "regnetx_040\n",
      "regnetx_064\n",
      "regnetx_080\n",
      "regnetx_120\n",
      "regnetx_160\n",
      "regnetx_320\n",
      "regnety_002\n",
      "regnety_004\n",
      "regnety_006\n",
      "regnety_008\n",
      "regnety_008_tv\n",
      "regnety_016\n",
      "regnety_032\n",
      "regnety_040\n",
      "regnety_040_sgn\n",
      "regnety_064\n",
      "regnety_080\n",
      "regnety_080_tv\n",
      "regnety_120\n",
      "regnety_160\n",
      "regnety_320\n",
      "regnety_640\n",
      "regnety_1280\n",
      "regnety_2560\n",
      "regnetz_005\n",
      "regnetz_040\n",
      "regnetz_040_h\n",
      "regnetz_b16\n",
      "regnetz_b16_evos\n",
      "regnetz_c16\n",
      "regnetz_c16_evos\n",
      "regnetz_d8\n",
      "regnetz_d8_evos\n",
      "regnetz_d32\n",
      "regnetz_e8\n",
      "repghostnet_050\n",
      "repghostnet_058\n",
      "repghostnet_080\n",
      "repghostnet_100\n",
      "repghostnet_111\n",
      "repghostnet_130\n",
      "repghostnet_150\n",
      "repghostnet_200\n",
      "repvgg_a0\n",
      "repvgg_a1\n",
      "repvgg_a2\n",
      "repvgg_b0\n",
      "repvgg_b1\n",
      "repvgg_b1g4\n",
      "repvgg_b2\n",
      "repvgg_b2g4\n",
      "repvgg_b3\n",
      "repvgg_b3g4\n",
      "repvgg_d2se\n",
      "repvit_m0_9\n",
      "repvit_m1\n",
      "repvit_m1_0\n",
      "repvit_m1_1\n",
      "repvit_m1_5\n",
      "repvit_m2\n",
      "repvit_m2_3\n",
      "repvit_m3\n",
      "res2net50_14w_8s\n",
      "res2net50_26w_4s\n",
      "res2net50_26w_6s\n",
      "res2net50_26w_8s\n",
      "res2net50_48w_2s\n",
      "res2net50d\n",
      "res2net101_26w_4s\n",
      "res2net101d\n",
      "res2next50\n",
      "resmlp_12_224\n",
      "resmlp_24_224\n",
      "resmlp_36_224\n",
      "resmlp_big_24_224\n",
      "resnest14d\n",
      "resnest26d\n",
      "resnest50d\n",
      "resnest50d_1s4x24d\n",
      "resnest50d_4s2x40d\n",
      "resnest101e\n",
      "resnest200e\n",
      "resnest269e\n",
      "resnet10t\n",
      "resnet14t\n",
      "resnet18\n",
      "resnet18d\n",
      "resnet26\n",
      "resnet26d\n",
      "resnet26t\n",
      "resnet32ts\n",
      "resnet33ts\n",
      "resnet34\n",
      "resnet34d\n",
      "resnet50\n",
      "resnet50_gn\n",
      "resnet50c\n",
      "resnet50d\n",
      "resnet50s\n",
      "resnet50t\n",
      "resnet51q\n",
      "resnet61q\n",
      "resnet101\n",
      "resnet101c\n",
      "resnet101d\n",
      "resnet101s\n",
      "resnet152\n",
      "resnet152c\n",
      "resnet152d\n",
      "resnet152s\n",
      "resnet200\n",
      "resnet200d\n",
      "resnetaa34d\n",
      "resnetaa50\n",
      "resnetaa50d\n",
      "resnetaa101d\n",
      "resnetblur18\n",
      "resnetblur50\n",
      "resnetblur50d\n",
      "resnetblur101d\n",
      "resnetrs50\n",
      "resnetrs101\n",
      "resnetrs152\n",
      "resnetrs200\n",
      "resnetrs270\n",
      "resnetrs350\n",
      "resnetrs420\n",
      "resnetv2_50\n",
      "resnetv2_50d\n",
      "resnetv2_50d_evos\n",
      "resnetv2_50d_frn\n",
      "resnetv2_50d_gn\n",
      "resnetv2_50t\n",
      "resnetv2_50x1_bit\n",
      "resnetv2_50x3_bit\n",
      "resnetv2_101\n",
      "resnetv2_101d\n",
      "resnetv2_101x1_bit\n",
      "resnetv2_101x3_bit\n",
      "resnetv2_152\n",
      "resnetv2_152d\n",
      "resnetv2_152x2_bit\n",
      "resnetv2_152x4_bit\n",
      "resnext26ts\n",
      "resnext50_32x4d\n",
      "resnext50d_32x4d\n",
      "resnext101_32x4d\n",
      "resnext101_32x8d\n",
      "resnext101_32x16d\n",
      "resnext101_32x32d\n",
      "resnext101_64x4d\n",
      "rexnet_100\n",
      "rexnet_130\n",
      "rexnet_150\n",
      "rexnet_200\n",
      "rexnet_300\n",
      "rexnetr_100\n",
      "rexnetr_130\n",
      "rexnetr_150\n",
      "rexnetr_200\n",
      "rexnetr_300\n",
      "samvit_base_patch16\n",
      "samvit_base_patch16_224\n",
      "samvit_huge_patch16\n",
      "samvit_large_patch16\n",
      "sebotnet33ts_256\n",
      "sedarknet21\n",
      "sehalonet33ts\n",
      "selecsls42\n",
      "selecsls42b\n",
      "selecsls60\n",
      "selecsls60b\n",
      "selecsls84\n",
      "semnasnet_050\n",
      "semnasnet_075\n",
      "semnasnet_100\n",
      "semnasnet_140\n",
      "senet154\n",
      "sequencer2d_l\n",
      "sequencer2d_m\n",
      "sequencer2d_s\n",
      "seresnet18\n",
      "seresnet33ts\n",
      "seresnet34\n",
      "seresnet50\n",
      "seresnet50t\n",
      "seresnet101\n",
      "seresnet152\n",
      "seresnet152d\n",
      "seresnet200d\n",
      "seresnet269d\n",
      "seresnetaa50d\n",
      "seresnext26d_32x4d\n",
      "seresnext26t_32x4d\n",
      "seresnext26ts\n",
      "seresnext50_32x4d\n",
      "seresnext101_32x4d\n",
      "seresnext101_32x8d\n",
      "seresnext101_64x4d\n",
      "seresnext101d_32x8d\n",
      "seresnextaa101d_32x8d\n",
      "seresnextaa201d_32x8d\n",
      "skresnet18\n",
      "skresnet34\n",
      "skresnet50\n",
      "skresnet50d\n",
      "skresnext50_32x4d\n",
      "spnasnet_100\n",
      "swin_base_patch4_window7_224\n",
      "swin_base_patch4_window12_384\n",
      "swin_large_patch4_window7_224\n",
      "swin_large_patch4_window12_384\n",
      "swin_s3_base_224\n",
      "swin_s3_small_224\n",
      "swin_s3_tiny_224\n",
      "swin_small_patch4_window7_224\n",
      "swin_tiny_patch4_window7_224\n",
      "swinv2_base_window8_256\n",
      "swinv2_base_window12_192\n",
      "swinv2_base_window12to16_192to256\n",
      "swinv2_base_window12to24_192to384\n",
      "swinv2_base_window16_256\n",
      "swinv2_cr_base_224\n",
      "swinv2_cr_base_384\n",
      "swinv2_cr_base_ns_224\n",
      "swinv2_cr_giant_224\n",
      "swinv2_cr_giant_384\n",
      "swinv2_cr_huge_224\n",
      "swinv2_cr_huge_384\n",
      "swinv2_cr_large_224\n",
      "swinv2_cr_large_384\n",
      "swinv2_cr_small_224\n",
      "swinv2_cr_small_384\n",
      "swinv2_cr_small_ns_224\n",
      "swinv2_cr_small_ns_256\n",
      "swinv2_cr_tiny_224\n",
      "swinv2_cr_tiny_384\n",
      "swinv2_cr_tiny_ns_224\n",
      "swinv2_large_window12_192\n",
      "swinv2_large_window12to16_192to256\n",
      "swinv2_large_window12to24_192to384\n",
      "swinv2_small_window8_256\n",
      "swinv2_small_window16_256\n",
      "swinv2_tiny_window8_256\n",
      "swinv2_tiny_window16_256\n",
      "tf_efficientnet_b0\n",
      "tf_efficientnet_b1\n",
      "tf_efficientnet_b2\n",
      "tf_efficientnet_b3\n",
      "tf_efficientnet_b4\n",
      "tf_efficientnet_b5\n",
      "tf_efficientnet_b6\n",
      "tf_efficientnet_b7\n",
      "tf_efficientnet_b8\n",
      "tf_efficientnet_cc_b0_4e\n",
      "tf_efficientnet_cc_b0_8e\n",
      "tf_efficientnet_cc_b1_8e\n",
      "tf_efficientnet_el\n",
      "tf_efficientnet_em\n",
      "tf_efficientnet_es\n",
      "tf_efficientnet_l2\n",
      "tf_efficientnet_lite0\n",
      "tf_efficientnet_lite1\n",
      "tf_efficientnet_lite2\n",
      "tf_efficientnet_lite3\n",
      "tf_efficientnet_lite4\n",
      "tf_efficientnetv2_b0\n",
      "tf_efficientnetv2_b1\n",
      "tf_efficientnetv2_b2\n",
      "tf_efficientnetv2_b3\n",
      "tf_efficientnetv2_l\n",
      "tf_efficientnetv2_m\n",
      "tf_efficientnetv2_s\n",
      "tf_efficientnetv2_xl\n",
      "tf_mixnet_l\n",
      "tf_mixnet_m\n",
      "tf_mixnet_s\n",
      "tf_mobilenetv3_large_075\n",
      "tf_mobilenetv3_large_100\n",
      "tf_mobilenetv3_large_minimal_100\n",
      "tf_mobilenetv3_small_075\n",
      "tf_mobilenetv3_small_100\n",
      "tf_mobilenetv3_small_minimal_100\n",
      "tiny_vit_5m_224\n",
      "tiny_vit_11m_224\n",
      "tiny_vit_21m_224\n",
      "tiny_vit_21m_384\n",
      "tiny_vit_21m_512\n",
      "tinynet_a\n",
      "tinynet_b\n",
      "tinynet_c\n",
      "tinynet_d\n",
      "tinynet_e\n",
      "tnt_b_patch16_224\n",
      "tnt_s_patch16_224\n",
      "tresnet_l\n",
      "tresnet_m\n",
      "tresnet_v2_l\n",
      "tresnet_xl\n",
      "twins_pcpvt_base\n",
      "twins_pcpvt_large\n",
      "twins_pcpvt_small\n",
      "twins_svt_base\n",
      "twins_svt_large\n",
      "twins_svt_small\n",
      "vgg11\n",
      "vgg11_bn\n",
      "vgg13\n",
      "vgg13_bn\n",
      "vgg16\n",
      "vgg16_bn\n",
      "vgg19\n",
      "vgg19_bn\n",
      "visformer_small\n",
      "visformer_tiny\n",
      "vit_base_patch8_224\n",
      "vit_base_patch14_dinov2\n",
      "vit_base_patch14_reg4_dinov2\n",
      "vit_base_patch16_18x2_224\n",
      "vit_base_patch16_224\n",
      "vit_base_patch16_224_miil\n",
      "vit_base_patch16_384\n",
      "vit_base_patch16_clip_224\n",
      "vit_base_patch16_clip_384\n",
      "vit_base_patch16_clip_quickgelu_224\n",
      "vit_base_patch16_gap_224\n",
      "vit_base_patch16_plus_240\n",
      "vit_base_patch16_reg4_gap_256\n",
      "vit_base_patch16_rpn_224\n",
      "vit_base_patch16_siglip_224\n",
      "vit_base_patch16_siglip_256\n",
      "vit_base_patch16_siglip_384\n",
      "vit_base_patch16_siglip_512\n",
      "vit_base_patch16_xp_224\n",
      "vit_base_patch32_224\n",
      "vit_base_patch32_384\n",
      "vit_base_patch32_clip_224\n",
      "vit_base_patch32_clip_256\n",
      "vit_base_patch32_clip_384\n",
      "vit_base_patch32_clip_448\n",
      "vit_base_patch32_clip_quickgelu_224\n",
      "vit_base_patch32_plus_256\n",
      "vit_base_r26_s32_224\n",
      "vit_base_r50_s16_224\n",
      "vit_base_r50_s16_384\n",
      "vit_base_resnet26d_224\n",
      "vit_base_resnet50d_224\n",
      "vit_giant_patch14_224\n",
      "vit_giant_patch14_clip_224\n",
      "vit_giant_patch14_dinov2\n",
      "vit_giant_patch14_reg4_dinov2\n",
      "vit_giant_patch16_gap_224\n",
      "vit_gigantic_patch14_224\n",
      "vit_gigantic_patch14_clip_224\n",
      "vit_huge_patch14_224\n",
      "vit_huge_patch14_clip_224\n",
      "vit_huge_patch14_clip_336\n",
      "vit_huge_patch14_clip_378\n",
      "vit_huge_patch14_clip_quickgelu_224\n",
      "vit_huge_patch14_clip_quickgelu_378\n",
      "vit_huge_patch14_gap_224\n",
      "vit_huge_patch14_xp_224\n",
      "vit_huge_patch16_gap_448\n",
      "vit_large_patch14_224\n",
      "vit_large_patch14_clip_224\n",
      "vit_large_patch14_clip_336\n",
      "vit_large_patch14_clip_quickgelu_224\n",
      "vit_large_patch14_clip_quickgelu_336\n",
      "vit_large_patch14_dinov2\n",
      "vit_large_patch14_reg4_dinov2\n",
      "vit_large_patch14_xp_224\n",
      "vit_large_patch16_224\n",
      "vit_large_patch16_384\n",
      "vit_large_patch16_siglip_256\n",
      "vit_large_patch16_siglip_384\n",
      "vit_large_patch32_224\n",
      "vit_large_patch32_384\n",
      "vit_large_r50_s32_224\n",
      "vit_large_r50_s32_384\n",
      "vit_medium_patch16_gap_240\n",
      "vit_medium_patch16_gap_256\n",
      "vit_medium_patch16_gap_384\n",
      "vit_medium_patch16_reg4_256\n",
      "vit_medium_patch16_reg4_gap_256\n",
      "vit_relpos_base_patch16_224\n",
      "vit_relpos_base_patch16_cls_224\n",
      "vit_relpos_base_patch16_clsgap_224\n",
      "vit_relpos_base_patch16_plus_240\n",
      "vit_relpos_base_patch16_rpn_224\n",
      "vit_relpos_base_patch32_plus_rpn_256\n",
      "vit_relpos_medium_patch16_224\n",
      "vit_relpos_medium_patch16_cls_224\n",
      "vit_relpos_medium_patch16_rpn_224\n",
      "vit_relpos_small_patch16_224\n",
      "vit_relpos_small_patch16_rpn_224\n",
      "vit_small_patch8_224\n",
      "vit_small_patch14_dinov2\n",
      "vit_small_patch14_reg4_dinov2\n",
      "vit_small_patch16_18x2_224\n",
      "vit_small_patch16_36x1_224\n",
      "vit_small_patch16_224\n",
      "vit_small_patch16_384\n",
      "vit_small_patch32_224\n",
      "vit_small_patch32_384\n",
      "vit_small_r26_s32_224\n",
      "vit_small_r26_s32_384\n",
      "vit_small_resnet26d_224\n",
      "vit_small_resnet50d_s16_224\n",
      "vit_so150m_patch16_reg4_gap_256\n",
      "vit_so150m_patch16_reg4_map_256\n",
      "vit_so400m_patch14_siglip_224\n",
      "vit_so400m_patch14_siglip_384\n",
      "vit_srelpos_medium_patch16_224\n",
      "vit_srelpos_small_patch16_224\n",
      "vit_tiny_patch16_224\n",
      "vit_tiny_patch16_384\n",
      "vit_tiny_r_s16_p8_224\n",
      "vit_tiny_r_s16_p8_384\n",
      "volo_d1_224\n",
      "volo_d1_384\n",
      "volo_d2_224\n",
      "volo_d2_384\n",
      "volo_d3_224\n",
      "volo_d3_448\n",
      "volo_d4_224\n",
      "volo_d4_448\n",
      "volo_d5_224\n",
      "volo_d5_448\n",
      "volo_d5_512\n",
      "vovnet39a\n",
      "vovnet57a\n",
      "wide_resnet50_2\n",
      "wide_resnet101_2\n",
      "xception41\n",
      "xception41p\n",
      "xception65\n",
      "xception65p\n",
      "xception71\n",
      "xcit_large_24_p8_224\n",
      "xcit_large_24_p8_384\n",
      "xcit_large_24_p16_224\n",
      "xcit_large_24_p16_384\n",
      "xcit_medium_24_p8_224\n",
      "xcit_medium_24_p8_384\n",
      "xcit_medium_24_p16_224\n",
      "xcit_medium_24_p16_384\n",
      "xcit_nano_12_p8_224\n",
      "xcit_nano_12_p8_384\n",
      "xcit_nano_12_p16_224\n",
      "xcit_nano_12_p16_384\n",
      "xcit_small_12_p8_224\n",
      "xcit_small_12_p8_384\n",
      "xcit_small_12_p16_224\n",
      "xcit_small_12_p16_384\n",
      "xcit_small_24_p8_224\n",
      "xcit_small_24_p8_384\n",
      "xcit_small_24_p16_224\n",
      "xcit_small_24_p16_384\n",
      "xcit_tiny_12_p8_224\n",
      "xcit_tiny_12_p8_384\n",
      "xcit_tiny_12_p16_224\n",
      "xcit_tiny_12_p16_384\n",
      "xcit_tiny_24_p8_224\n",
      "xcit_tiny_24_p8_384\n",
      "xcit_tiny_24_p16_224\n",
      "xcit_tiny_24_p16_384\n"
     ]
    }
   ],
   "source": [
    "import timm\n",
    "for m in timm.list_models(): print(m)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "resnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==========================================================================================\n",
       "Layer (type:depth-idx)                   Output Shape              Param #\n",
       "==========================================================================================\n",
       "ResNet                                   [1, 10]                   --\n",
       "├─Conv2d: 1-1                            [1, 64, 112, 112]         9,408\n",
       "├─BatchNorm2d: 1-2                       [1, 64, 112, 112]         128\n",
       "├─ReLU: 1-3                              [1, 64, 112, 112]         --\n",
       "├─MaxPool2d: 1-4                         [1, 64, 56, 56]           --\n",
       "├─Sequential: 1-5                        [1, 256, 56, 56]          --\n",
       "│    └─Bottleneck: 2-1                   [1, 256, 56, 56]          --\n",
       "│    │    └─Conv2d: 3-1                  [1, 64, 56, 56]           4,096\n",
       "│    │    └─BatchNorm2d: 3-2             [1, 64, 56, 56]           128\n",
       "│    │    └─ReLU: 3-3                    [1, 64, 56, 56]           --\n",
       "│    │    └─Conv2d: 3-4                  [1, 64, 56, 56]           36,864\n",
       "│    │    └─BatchNorm2d: 3-5             [1, 64, 56, 56]           128\n",
       "│    │    └─Identity: 3-6                [1, 64, 56, 56]           --\n",
       "│    │    └─ReLU: 3-7                    [1, 64, 56, 56]           --\n",
       "│    │    └─Identity: 3-8                [1, 64, 56, 56]           --\n",
       "│    │    └─Conv2d: 3-9                  [1, 256, 56, 56]          16,384\n",
       "│    │    └─BatchNorm2d: 3-10            [1, 256, 56, 56]          512\n",
       "│    │    └─Sequential: 3-11             [1, 256, 56, 56]          16,896\n",
       "│    │    └─ReLU: 3-12                   [1, 256, 56, 56]          --\n",
       "│    └─Bottleneck: 2-2                   [1, 256, 56, 56]          --\n",
       "│    │    └─Conv2d: 3-13                 [1, 64, 56, 56]           16,384\n",
       "│    │    └─BatchNorm2d: 3-14            [1, 64, 56, 56]           128\n",
       "│    │    └─ReLU: 3-15                   [1, 64, 56, 56]           --\n",
       "│    │    └─Conv2d: 3-16                 [1, 64, 56, 56]           36,864\n",
       "│    │    └─BatchNorm2d: 3-17            [1, 64, 56, 56]           128\n",
       "│    │    └─Identity: 3-18               [1, 64, 56, 56]           --\n",
       "│    │    └─ReLU: 3-19                   [1, 64, 56, 56]           --\n",
       "│    │    └─Identity: 3-20               [1, 64, 56, 56]           --\n",
       "│    │    └─Conv2d: 3-21                 [1, 256, 56, 56]          16,384\n",
       "│    │    └─BatchNorm2d: 3-22            [1, 256, 56, 56]          512\n",
       "│    │    └─ReLU: 3-23                   [1, 256, 56, 56]          --\n",
       "│    └─Bottleneck: 2-3                   [1, 256, 56, 56]          --\n",
       "│    │    └─Conv2d: 3-24                 [1, 64, 56, 56]           16,384\n",
       "│    │    └─BatchNorm2d: 3-25            [1, 64, 56, 56]           128\n",
       "│    │    └─ReLU: 3-26                   [1, 64, 56, 56]           --\n",
       "│    │    └─Conv2d: 3-27                 [1, 64, 56, 56]           36,864\n",
       "│    │    └─BatchNorm2d: 3-28            [1, 64, 56, 56]           128\n",
       "│    │    └─Identity: 3-29               [1, 64, 56, 56]           --\n",
       "│    │    └─ReLU: 3-30                   [1, 64, 56, 56]           --\n",
       "│    │    └─Identity: 3-31               [1, 64, 56, 56]           --\n",
       "│    │    └─Conv2d: 3-32                 [1, 256, 56, 56]          16,384\n",
       "│    │    └─BatchNorm2d: 3-33            [1, 256, 56, 56]          512\n",
       "│    │    └─ReLU: 3-34                   [1, 256, 56, 56]          --\n",
       "├─Sequential: 1-6                        [1, 512, 28, 28]          --\n",
       "│    └─Bottleneck: 2-4                   [1, 512, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-35                 [1, 128, 56, 56]          32,768\n",
       "│    │    └─BatchNorm2d: 3-36            [1, 128, 56, 56]          256\n",
       "│    │    └─ReLU: 3-37                   [1, 128, 56, 56]          --\n",
       "│    │    └─Conv2d: 3-38                 [1, 128, 28, 28]          147,456\n",
       "│    │    └─BatchNorm2d: 3-39            [1, 128, 28, 28]          256\n",
       "│    │    └─Identity: 3-40               [1, 128, 28, 28]          --\n",
       "│    │    └─ReLU: 3-41                   [1, 128, 28, 28]          --\n",
       "│    │    └─Identity: 3-42               [1, 128, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-43                 [1, 512, 28, 28]          65,536\n",
       "│    │    └─BatchNorm2d: 3-44            [1, 512, 28, 28]          1,024\n",
       "│    │    └─Sequential: 3-45             [1, 512, 28, 28]          132,096\n",
       "│    │    └─ReLU: 3-46                   [1, 512, 28, 28]          --\n",
       "│    └─Bottleneck: 2-5                   [1, 512, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-47                 [1, 128, 28, 28]          65,536\n",
       "│    │    └─BatchNorm2d: 3-48            [1, 128, 28, 28]          256\n",
       "│    │    └─ReLU: 3-49                   [1, 128, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-50                 [1, 128, 28, 28]          147,456\n",
       "│    │    └─BatchNorm2d: 3-51            [1, 128, 28, 28]          256\n",
       "│    │    └─Identity: 3-52               [1, 128, 28, 28]          --\n",
       "│    │    └─ReLU: 3-53                   [1, 128, 28, 28]          --\n",
       "│    │    └─Identity: 3-54               [1, 128, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-55                 [1, 512, 28, 28]          65,536\n",
       "│    │    └─BatchNorm2d: 3-56            [1, 512, 28, 28]          1,024\n",
       "│    │    └─ReLU: 3-57                   [1, 512, 28, 28]          --\n",
       "│    └─Bottleneck: 2-6                   [1, 512, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-58                 [1, 128, 28, 28]          65,536\n",
       "│    │    └─BatchNorm2d: 3-59            [1, 128, 28, 28]          256\n",
       "│    │    └─ReLU: 3-60                   [1, 128, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-61                 [1, 128, 28, 28]          147,456\n",
       "│    │    └─BatchNorm2d: 3-62            [1, 128, 28, 28]          256\n",
       "│    │    └─Identity: 3-63               [1, 128, 28, 28]          --\n",
       "│    │    └─ReLU: 3-64                   [1, 128, 28, 28]          --\n",
       "│    │    └─Identity: 3-65               [1, 128, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-66                 [1, 512, 28, 28]          65,536\n",
       "│    │    └─BatchNorm2d: 3-67            [1, 512, 28, 28]          1,024\n",
       "│    │    └─ReLU: 3-68                   [1, 512, 28, 28]          --\n",
       "│    └─Bottleneck: 2-7                   [1, 512, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-69                 [1, 128, 28, 28]          65,536\n",
       "│    │    └─BatchNorm2d: 3-70            [1, 128, 28, 28]          256\n",
       "│    │    └─ReLU: 3-71                   [1, 128, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-72                 [1, 128, 28, 28]          147,456\n",
       "│    │    └─BatchNorm2d: 3-73            [1, 128, 28, 28]          256\n",
       "│    │    └─Identity: 3-74               [1, 128, 28, 28]          --\n",
       "│    │    └─ReLU: 3-75                   [1, 128, 28, 28]          --\n",
       "│    │    └─Identity: 3-76               [1, 128, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-77                 [1, 512, 28, 28]          65,536\n",
       "│    │    └─BatchNorm2d: 3-78            [1, 512, 28, 28]          1,024\n",
       "│    │    └─ReLU: 3-79                   [1, 512, 28, 28]          --\n",
       "├─Sequential: 1-7                        [1, 1024, 14, 14]         --\n",
       "│    └─Bottleneck: 2-8                   [1, 1024, 14, 14]         --\n",
       "│    │    └─Conv2d: 3-80                 [1, 256, 28, 28]          131,072\n",
       "│    │    └─BatchNorm2d: 3-81            [1, 256, 28, 28]          512\n",
       "│    │    └─ReLU: 3-82                   [1, 256, 28, 28]          --\n",
       "│    │    └─Conv2d: 3-83                 [1, 256, 14, 14]          589,824\n",
       "│    │    └─BatchNorm2d: 3-84            [1, 256, 14, 14]          512\n",
       "│    │    └─Identity: 3-85               [1, 256, 14, 14]          --\n",
       "│    │    └─ReLU: 3-86                   [1, 256, 14, 14]          --\n",
       "│    │    └─Identity: 3-87               [1, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-88                 [1, 1024, 14, 14]         262,144\n",
       "│    │    └─BatchNorm2d: 3-89            [1, 1024, 14, 14]         2,048\n",
       "│    │    └─Sequential: 3-90             [1, 1024, 14, 14]         526,336\n",
       "│    │    └─ReLU: 3-91                   [1, 1024, 14, 14]         --\n",
       "│    └─Bottleneck: 2-9                   [1, 1024, 14, 14]         --\n",
       "│    │    └─Conv2d: 3-92                 [1, 256, 14, 14]          262,144\n",
       "│    │    └─BatchNorm2d: 3-93            [1, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-94                   [1, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-95                 [1, 256, 14, 14]          589,824\n",
       "│    │    └─BatchNorm2d: 3-96            [1, 256, 14, 14]          512\n",
       "│    │    └─Identity: 3-97               [1, 256, 14, 14]          --\n",
       "│    │    └─ReLU: 3-98                   [1, 256, 14, 14]          --\n",
       "│    │    └─Identity: 3-99               [1, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-100                [1, 1024, 14, 14]         262,144\n",
       "│    │    └─BatchNorm2d: 3-101           [1, 1024, 14, 14]         2,048\n",
       "│    │    └─ReLU: 3-102                  [1, 1024, 14, 14]         --\n",
       "│    └─Bottleneck: 2-10                  [1, 1024, 14, 14]         --\n",
       "│    │    └─Conv2d: 3-103                [1, 256, 14, 14]          262,144\n",
       "│    │    └─BatchNorm2d: 3-104           [1, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-105                  [1, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-106                [1, 256, 14, 14]          589,824\n",
       "│    │    └─BatchNorm2d: 3-107           [1, 256, 14, 14]          512\n",
       "│    │    └─Identity: 3-108              [1, 256, 14, 14]          --\n",
       "│    │    └─ReLU: 3-109                  [1, 256, 14, 14]          --\n",
       "│    │    └─Identity: 3-110              [1, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-111                [1, 1024, 14, 14]         262,144\n",
       "│    │    └─BatchNorm2d: 3-112           [1, 1024, 14, 14]         2,048\n",
       "│    │    └─ReLU: 3-113                  [1, 1024, 14, 14]         --\n",
       "│    └─Bottleneck: 2-11                  [1, 1024, 14, 14]         --\n",
       "│    │    └─Conv2d: 3-114                [1, 256, 14, 14]          262,144\n",
       "│    │    └─BatchNorm2d: 3-115           [1, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-116                  [1, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-117                [1, 256, 14, 14]          589,824\n",
       "│    │    └─BatchNorm2d: 3-118           [1, 256, 14, 14]          512\n",
       "│    │    └─Identity: 3-119              [1, 256, 14, 14]          --\n",
       "│    │    └─ReLU: 3-120                  [1, 256, 14, 14]          --\n",
       "│    │    └─Identity: 3-121              [1, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-122                [1, 1024, 14, 14]         262,144\n",
       "│    │    └─BatchNorm2d: 3-123           [1, 1024, 14, 14]         2,048\n",
       "│    │    └─ReLU: 3-124                  [1, 1024, 14, 14]         --\n",
       "│    └─Bottleneck: 2-12                  [1, 1024, 14, 14]         --\n",
       "│    │    └─Conv2d: 3-125                [1, 256, 14, 14]          262,144\n",
       "│    │    └─BatchNorm2d: 3-126           [1, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-127                  [1, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-128                [1, 256, 14, 14]          589,824\n",
       "│    │    └─BatchNorm2d: 3-129           [1, 256, 14, 14]          512\n",
       "│    │    └─Identity: 3-130              [1, 256, 14, 14]          --\n",
       "│    │    └─ReLU: 3-131                  [1, 256, 14, 14]          --\n",
       "│    │    └─Identity: 3-132              [1, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-133                [1, 1024, 14, 14]         262,144\n",
       "│    │    └─BatchNorm2d: 3-134           [1, 1024, 14, 14]         2,048\n",
       "│    │    └─ReLU: 3-135                  [1, 1024, 14, 14]         --\n",
       "│    └─Bottleneck: 2-13                  [1, 1024, 14, 14]         --\n",
       "│    │    └─Conv2d: 3-136                [1, 256, 14, 14]          262,144\n",
       "│    │    └─BatchNorm2d: 3-137           [1, 256, 14, 14]          512\n",
       "│    │    └─ReLU: 3-138                  [1, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-139                [1, 256, 14, 14]          589,824\n",
       "│    │    └─BatchNorm2d: 3-140           [1, 256, 14, 14]          512\n",
       "│    │    └─Identity: 3-141              [1, 256, 14, 14]          --\n",
       "│    │    └─ReLU: 3-142                  [1, 256, 14, 14]          --\n",
       "│    │    └─Identity: 3-143              [1, 256, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-144                [1, 1024, 14, 14]         262,144\n",
       "│    │    └─BatchNorm2d: 3-145           [1, 1024, 14, 14]         2,048\n",
       "│    │    └─ReLU: 3-146                  [1, 1024, 14, 14]         --\n",
       "├─Sequential: 1-8                        [1, 2048, 7, 7]           --\n",
       "│    └─Bottleneck: 2-14                  [1, 2048, 7, 7]           --\n",
       "│    │    └─Conv2d: 3-147                [1, 512, 14, 14]          524,288\n",
       "│    │    └─BatchNorm2d: 3-148           [1, 512, 14, 14]          1,024\n",
       "│    │    └─ReLU: 3-149                  [1, 512, 14, 14]          --\n",
       "│    │    └─Conv2d: 3-150                [1, 512, 7, 7]            2,359,296\n",
       "│    │    └─BatchNorm2d: 3-151           [1, 512, 7, 7]            1,024\n",
       "│    │    └─Identity: 3-152              [1, 512, 7, 7]            --\n",
       "│    │    └─ReLU: 3-153                  [1, 512, 7, 7]            --\n",
       "│    │    └─Identity: 3-154              [1, 512, 7, 7]            --\n",
       "│    │    └─Conv2d: 3-155                [1, 2048, 7, 7]           1,048,576\n",
       "│    │    └─BatchNorm2d: 3-156           [1, 2048, 7, 7]           4,096\n",
       "│    │    └─Sequential: 3-157            [1, 2048, 7, 7]           2,101,248\n",
       "│    │    └─ReLU: 3-158                  [1, 2048, 7, 7]           --\n",
       "│    └─Bottleneck: 2-15                  [1, 2048, 7, 7]           --\n",
       "│    │    └─Conv2d: 3-159                [1, 512, 7, 7]            1,048,576\n",
       "│    │    └─BatchNorm2d: 3-160           [1, 512, 7, 7]            1,024\n",
       "│    │    └─ReLU: 3-161                  [1, 512, 7, 7]            --\n",
       "│    │    └─Conv2d: 3-162                [1, 512, 7, 7]            2,359,296\n",
       "│    │    └─BatchNorm2d: 3-163           [1, 512, 7, 7]            1,024\n",
       "│    │    └─Identity: 3-164              [1, 512, 7, 7]            --\n",
       "│    │    └─ReLU: 3-165                  [1, 512, 7, 7]            --\n",
       "│    │    └─Identity: 3-166              [1, 512, 7, 7]            --\n",
       "│    │    └─Conv2d: 3-167                [1, 2048, 7, 7]           1,048,576\n",
       "│    │    └─BatchNorm2d: 3-168           [1, 2048, 7, 7]           4,096\n",
       "│    │    └─ReLU: 3-169                  [1, 2048, 7, 7]           --\n",
       "│    └─Bottleneck: 2-16                  [1, 2048, 7, 7]           --\n",
       "│    │    └─Conv2d: 3-170                [1, 512, 7, 7]            1,048,576\n",
       "│    │    └─BatchNorm2d: 3-171           [1, 512, 7, 7]            1,024\n",
       "│    │    └─ReLU: 3-172                  [1, 512, 7, 7]            --\n",
       "│    │    └─Conv2d: 3-173                [1, 512, 7, 7]            2,359,296\n",
       "│    │    └─BatchNorm2d: 3-174           [1, 512, 7, 7]            1,024\n",
       "│    │    └─Identity: 3-175              [1, 512, 7, 7]            --\n",
       "│    │    └─ReLU: 3-176                  [1, 512, 7, 7]            --\n",
       "│    │    └─Identity: 3-177              [1, 512, 7, 7]            --\n",
       "│    │    └─Conv2d: 3-178                [1, 2048, 7, 7]           1,048,576\n",
       "│    │    └─BatchNorm2d: 3-179           [1, 2048, 7, 7]           4,096\n",
       "│    │    └─ReLU: 3-180                  [1, 2048, 7, 7]           --\n",
       "├─SelectAdaptivePool2d: 1-9              [1, 2048]                 --\n",
       "│    └─AdaptiveMaxPool2d: 2-17           [1, 2048, 1, 1]           --\n",
       "│    └─Flatten: 2-18                     [1, 2048]                 --\n",
       "├─Linear: 1-10                           [1, 10]                   20,490\n",
       "==========================================================================================\n",
       "Total params: 23,528,522\n",
       "Trainable params: 23,528,522\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.GIGABYTES): 4.09\n",
       "==========================================================================================\n",
       "Input size (MB): 0.60\n",
       "Forward/backward pass size (MB): 177.82\n",
       "Params size (MB): 94.11\n",
       "Estimated Total Size (MB): 272.54\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import timm\n",
    "from torchinfo import summary\n",
    "\n",
    "net = timm.create_model(\n",
    "    'resnet50', \n",
    "    pretrained=False,       # no pretrained weights \n",
    "    num_classes=10,         # output classes 10\n",
    "    in_chans=3,             # input channels \n",
    "    drop_rate=0.5,          # dropout 0.5\n",
    "    global_pool='max'       # pooling max\n",
    ")\n",
    "\n",
    "summary(net,(1,3,224,224)) # bs, in_channel, image size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "fine-tune protocol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import timm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "# 1. load pretrained model\n",
    "model = timm.create_model('resnet50', pretrained=True, num_classes=10)\n",
    "\n",
    "# 2. freeze pretrained model parameters\n",
    "for param in model.parameters():\n",
    "    param.requires_grad = False # will not update the weights of pretrained model\n",
    "\n",
    "for name, param in model.named_parameters():\n",
    "    if 'layer4' not in name: # freeze all layers except layer4\n",
    "        param.requires_grad = False\n",
    "\n",
    "# 3. only train the last layer\n",
    "# unfreeze fully connected layer parameters\n",
    "for param in model.fc.parameters():\n",
    "    param.requires_grad = True\n",
    "\n",
    "# 4. loss function and optimizer\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.fc.parameters(), lr=1e-4)  # only update the classification layer parameters \n",
    "\n",
    "# 5. prepare dataset and dataloader\n",
    "\n",
    "# 6. start training\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## models from MONAI\n",
    "* Reference: https://docs.monai.io/en/stable/networks.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Basic unet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BasicUNet features: (16, 32, 64, 128, 256, 16).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "====================================================================================================\n",
       "Layer (type:depth-idx)                             Output Shape              Param #\n",
       "====================================================================================================\n",
       "BasicUNet                                          [1, 1, 224]               --\n",
       "├─TwoConv: 1-1                                     [1, 16, 224]              --\n",
       "│    └─Convolution: 2-1                            [1, 16, 224]              --\n",
       "│    │    └─Conv1d: 3-1                            [1, 16, 224]              112\n",
       "│    │    └─ADN: 3-2                               [1, 16, 224]              32\n",
       "│    └─Convolution: 2-2                            [1, 16, 224]              --\n",
       "│    │    └─Conv1d: 3-3                            [1, 16, 224]              784\n",
       "│    │    └─ADN: 3-4                               [1, 16, 224]              32\n",
       "├─Down: 1-2                                        [1, 32, 112]              --\n",
       "│    └─MaxPool1d: 2-3                              [1, 16, 112]              --\n",
       "│    └─TwoConv: 2-4                                [1, 32, 112]              --\n",
       "│    │    └─Convolution: 3-5                       [1, 32, 112]              1,632\n",
       "│    │    └─Convolution: 3-6                       [1, 32, 112]              3,168\n",
       "├─Down: 1-3                                        [1, 64, 56]               --\n",
       "│    └─MaxPool1d: 2-5                              [1, 32, 56]               --\n",
       "│    └─TwoConv: 2-6                                [1, 64, 56]               --\n",
       "│    │    └─Convolution: 3-7                       [1, 64, 56]               6,336\n",
       "│    │    └─Convolution: 3-8                       [1, 64, 56]               12,480\n",
       "├─Down: 1-4                                        [1, 128, 28]              --\n",
       "│    └─MaxPool1d: 2-7                              [1, 64, 28]               --\n",
       "│    └─TwoConv: 2-8                                [1, 128, 28]              --\n",
       "│    │    └─Convolution: 3-9                       [1, 128, 28]              24,960\n",
       "│    │    └─Convolution: 3-10                      [1, 128, 28]              49,536\n",
       "├─Down: 1-5                                        [1, 256, 14]              --\n",
       "│    └─MaxPool1d: 2-9                              [1, 128, 14]              --\n",
       "│    └─TwoConv: 2-10                               [1, 256, 14]              --\n",
       "│    │    └─Convolution: 3-11                      [1, 256, 14]              99,072\n",
       "│    │    └─Convolution: 3-12                      [1, 256, 14]              197,376\n",
       "├─UpCat: 1-6                                       [1, 128, 28]              --\n",
       "│    └─UpSample: 2-11                              [1, 128, 28]              --\n",
       "│    │    └─ConvTranspose1d: 3-13                  [1, 128, 28]              65,664\n",
       "│    └─TwoConv: 2-12                               [1, 128, 28]              --\n",
       "│    │    └─Convolution: 3-14                      [1, 128, 28]              98,688\n",
       "│    │    └─Convolution: 3-15                      [1, 128, 28]              49,536\n",
       "├─UpCat: 1-7                                       [1, 64, 56]               --\n",
       "│    └─UpSample: 2-13                              [1, 64, 56]               --\n",
       "│    │    └─ConvTranspose1d: 3-16                  [1, 64, 56]               16,448\n",
       "│    └─TwoConv: 2-14                               [1, 64, 56]               --\n",
       "│    │    └─Convolution: 3-17                      [1, 64, 56]               24,768\n",
       "│    │    └─Convolution: 3-18                      [1, 64, 56]               12,480\n",
       "├─UpCat: 1-8                                       [1, 32, 112]              --\n",
       "│    └─UpSample: 2-15                              [1, 32, 112]              --\n",
       "│    │    └─ConvTranspose1d: 3-19                  [1, 32, 112]              4,128\n",
       "│    └─TwoConv: 2-16                               [1, 32, 112]              --\n",
       "│    │    └─Convolution: 3-20                      [1, 32, 112]              6,240\n",
       "│    │    └─Convolution: 3-21                      [1, 32, 112]              3,168\n",
       "├─UpCat: 1-9                                       [1, 16, 224]              --\n",
       "│    └─UpSample: 2-17                              [1, 32, 224]              --\n",
       "│    │    └─ConvTranspose1d: 3-22                  [1, 32, 224]              2,080\n",
       "│    └─TwoConv: 2-18                               [1, 16, 224]              --\n",
       "│    │    └─Convolution: 3-23                      [1, 16, 224]              2,352\n",
       "│    │    └─Convolution: 3-24                      [1, 16, 224]              816\n",
       "├─Conv1d: 1-10                                     [1, 1, 224]               17\n",
       "====================================================================================================\n",
       "Total params: 681,905\n",
       "Trainable params: 681,905\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 19.61\n",
       "====================================================================================================\n",
       "Input size (MB): 0.00\n",
       "Forward/backward pass size (MB): 1.18\n",
       "Params size (MB): 2.73\n",
       "Estimated Total Size (MB): 3.91\n",
       "===================================================================================================="
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from monai.networks.nets import BasicUnet\n",
    "from torchinfo import summary\n",
    "\n",
    "net = BasicUnet(\n",
    "    spatial_dims=1, \n",
    "    in_channels=2, \n",
    "    out_channels=1, \n",
    "    features=(16, 32, 64, 128, 256, 16), \n",
    "    act=('LeakyReLU', {'inplace': True, 'negative_slope': 0.1}), \n",
    "    norm=('instance', {'affine': True}), \n",
    "    bias=True, \n",
    "    dropout=0.0, \n",
    "    upsample='deconv'\n",
    "    )\n",
    "\n",
    "summary(net,(1,2,224)) # bs, in_channel, vecter length"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "v-net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "====================================================================================================\n",
       "Layer (type:depth-idx)                             Output Shape              Param #\n",
       "====================================================================================================\n",
       "VNet                                               [1, 1, 224, 224]          --\n",
       "├─InputTransition: 1-1                             [1, 16, 224, 224]         --\n",
       "│    └─Convolution: 2-1                            [1, 16, 224, 224]         --\n",
       "│    │    └─Conv2d: 3-1                            [1, 16, 224, 224]         800\n",
       "│    │    └─ADN: 3-2                               [1, 16, 224, 224]         32\n",
       "│    └─ELU: 2-2                                    [1, 16, 224, 224]         --\n",
       "├─DownTransition: 1-2                              [1, 32, 112, 112]         --\n",
       "│    └─Conv2d: 2-3                                 [1, 32, 112, 112]         2,048\n",
       "│    └─BatchNorm2d: 2-4                            [1, 32, 112, 112]         64\n",
       "│    └─ELU: 2-5                                    [1, 32, 112, 112]         --\n",
       "│    └─Sequential: 2-6                             [1, 32, 112, 112]         --\n",
       "│    │    └─LUConv: 3-3                            [1, 32, 112, 112]         25,664\n",
       "│    └─ELU: 2-7                                    [1, 32, 112, 112]         --\n",
       "├─DownTransition: 1-3                              [1, 64, 56, 56]           --\n",
       "│    └─Conv2d: 2-8                                 [1, 64, 56, 56]           8,192\n",
       "│    └─BatchNorm2d: 2-9                            [1, 64, 56, 56]           128\n",
       "│    └─ELU: 2-10                                   [1, 64, 56, 56]           --\n",
       "│    └─Sequential: 2-11                            [1, 64, 56, 56]           --\n",
       "│    │    └─LUConv: 3-4                            [1, 64, 56, 56]           102,528\n",
       "│    │    └─LUConv: 3-5                            [1, 64, 56, 56]           102,528\n",
       "│    └─ELU: 2-12                                   [1, 64, 56, 56]           --\n",
       "├─DownTransition: 1-4                              [1, 128, 28, 28]          --\n",
       "│    └─Conv2d: 2-13                                [1, 128, 28, 28]          32,768\n",
       "│    └─BatchNorm2d: 2-14                           [1, 128, 28, 28]          256\n",
       "│    └─ELU: 2-15                                   [1, 128, 28, 28]          --\n",
       "│    └─Dropout3d: 2-16                             [1, 128, 28, 28]          --\n",
       "│    └─Sequential: 2-17                            [1, 128, 28, 28]          --\n",
       "│    │    └─LUConv: 3-6                            [1, 128, 28, 28]          409,856\n",
       "│    │    └─LUConv: 3-7                            [1, 128, 28, 28]          409,856\n",
       "│    │    └─LUConv: 3-8                            [1, 128, 28, 28]          409,856\n",
       "│    └─ELU: 2-18                                   [1, 128, 28, 28]          --\n",
       "├─DownTransition: 1-5                              [1, 256, 14, 14]          --\n",
       "│    └─Conv2d: 2-19                                [1, 256, 14, 14]          131,072\n",
       "│    └─BatchNorm2d: 2-20                           [1, 256, 14, 14]          512\n",
       "│    └─ELU: 2-21                                   [1, 256, 14, 14]          --\n",
       "│    └─Dropout3d: 2-22                             [1, 256, 14, 14]          --\n",
       "│    └─Sequential: 2-23                            [1, 256, 14, 14]          --\n",
       "│    │    └─LUConv: 3-9                            [1, 256, 14, 14]          1,638,912\n",
       "│    │    └─LUConv: 3-10                           [1, 256, 14, 14]          1,638,912\n",
       "│    └─ELU: 2-24                                   [1, 256, 14, 14]          --\n",
       "├─UpTransition: 1-6                                [1, 256, 28, 28]          --\n",
       "│    └─Dropout3d: 2-25                             [1, 256, 14, 14]          --\n",
       "│    └─Dropout3d: 2-26                             [1, 128, 28, 28]          --\n",
       "│    └─ConvTranspose2d: 2-27                       [1, 128, 28, 28]          131,200\n",
       "│    └─BatchNorm2d: 2-28                           [1, 128, 28, 28]          256\n",
       "│    └─ELU: 2-29                                   [1, 128, 28, 28]          --\n",
       "│    └─Sequential: 2-30                            [1, 256, 28, 28]          --\n",
       "│    │    └─LUConv: 3-11                           [1, 256, 28, 28]          1,638,912\n",
       "│    │    └─LUConv: 3-12                           [1, 256, 28, 28]          1,638,912\n",
       "│    └─ELU: 2-31                                   [1, 256, 28, 28]          --\n",
       "├─UpTransition: 1-7                                [1, 128, 56, 56]          --\n",
       "│    └─Dropout3d: 2-32                             [1, 256, 28, 28]          --\n",
       "│    └─Dropout3d: 2-33                             [1, 64, 56, 56]           --\n",
       "│    └─ConvTranspose2d: 2-34                       [1, 64, 56, 56]           65,600\n",
       "│    └─BatchNorm2d: 2-35                           [1, 64, 56, 56]           128\n",
       "│    └─ELU: 2-36                                   [1, 64, 56, 56]           --\n",
       "│    └─Sequential: 2-37                            [1, 128, 56, 56]          --\n",
       "│    │    └─LUConv: 3-13                           [1, 128, 56, 56]          409,856\n",
       "│    │    └─LUConv: 3-14                           [1, 128, 56, 56]          409,856\n",
       "│    └─ELU: 2-38                                   [1, 128, 56, 56]          --\n",
       "├─UpTransition: 1-8                                [1, 64, 112, 112]         --\n",
       "│    └─Dropout3d: 2-39                             [1, 32, 112, 112]         --\n",
       "│    └─ConvTranspose2d: 2-40                       [1, 32, 112, 112]         16,416\n",
       "│    └─BatchNorm2d: 2-41                           [1, 32, 112, 112]         64\n",
       "│    └─ELU: 2-42                                   [1, 32, 112, 112]         --\n",
       "│    └─Sequential: 2-43                            [1, 64, 112, 112]         --\n",
       "│    │    └─LUConv: 3-15                           [1, 64, 112, 112]         102,528\n",
       "│    └─ELU: 2-44                                   [1, 64, 112, 112]         --\n",
       "├─UpTransition: 1-9                                [1, 32, 224, 224]         --\n",
       "│    └─Dropout3d: 2-45                             [1, 16, 224, 224]         --\n",
       "│    └─ConvTranspose2d: 2-46                       [1, 16, 224, 224]         4,112\n",
       "│    └─BatchNorm2d: 2-47                           [1, 16, 224, 224]         32\n",
       "│    └─ELU: 2-48                                   [1, 16, 224, 224]         --\n",
       "│    └─Sequential: 2-49                            [1, 32, 224, 224]         --\n",
       "│    │    └─LUConv: 3-16                           [1, 32, 224, 224]         25,664\n",
       "│    └─ELU: 2-50                                   [1, 32, 224, 224]         --\n",
       "├─OutputTransition: 1-10                           [1, 1, 224, 224]          --\n",
       "│    └─Convolution: 2-51                           [1, 1, 224, 224]          --\n",
       "│    │    └─Conv2d: 3-17                           [1, 1, 224, 224]          800\n",
       "│    │    └─ADN: 3-18                              [1, 1, 224, 224]          2\n",
       "│    └─ELU: 2-52                                   [1, 1, 224, 224]          --\n",
       "│    └─Conv2d: 2-53                                [1, 1, 224, 224]          2\n",
       "====================================================================================================\n",
       "Total params: 9,358,324\n",
       "Trainable params: 9,358,324\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.GIGABYTES): 11.18\n",
       "====================================================================================================\n",
       "Input size (MB): 0.40\n",
       "Forward/backward pass size (MB): 127.25\n",
       "Params size (MB): 37.43\n",
       "Estimated Total Size (MB): 165.08\n",
       "===================================================================================================="
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from monai.networks.nets import VNet\n",
    "from torchinfo import summary\n",
    "\n",
    "\n",
    "net = VNet(\n",
    "    spatial_dims=2,\n",
    "    in_channels=2, \n",
    "    out_channels=1, \n",
    "    act=('elu', {'inplace': True}), \n",
    "    dropout_prob=0.5, \n",
    "    dropout_prob_down=0.5, \n",
    "    dropout_prob_up=(0.5, 0.5), \n",
    "    dropout_dim=3, \n",
    "    bias=False\n",
    "   )\n",
    "\n",
    "summary(net,(1,2,224,224)) # bs, in_channel, vecter length"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "residual unet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "============================================================================================================================================\n",
       "Layer (type:depth-idx)                                                                     Output Shape              Param #\n",
       "============================================================================================================================================\n",
       "UNet                                                                                       [1, 1, 224]               --\n",
       "├─Sequential: 1-1                                                                          [1, 1, 224]               --\n",
       "│    └─ResidualUnit: 2-1                                                                   [1, 16, 112]              --\n",
       "│    │    └─Conv1d: 3-1                                                                    [1, 16, 112]              112\n",
       "│    │    └─Sequential: 3-2                                                                [1, 16, 112]              2,468\n",
       "│    └─SkipConnection: 2-2                                                                 [1, 32, 112]              --\n",
       "│    │    └─Sequential: 3-3                                                                [1, 16, 112]              1,086,806\n",
       "│    └─Sequential: 2-3                                                                     [1, 1, 224]               --\n",
       "│    │    └─Convolution: 3-4                                                               [1, 1, 224]               98\n",
       "│    │    └─ResidualUnit: 3-5                                                              [1, 1, 224]               4\n",
       "============================================================================================================================================\n",
       "Total params: 1,089,488\n",
       "Trainable params: 1,089,488\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 18.97\n",
       "============================================================================================================================================\n",
       "Input size (MB): 0.00\n",
       "Forward/backward pass size (MB): 0.95\n",
       "Params size (MB): 4.36\n",
       "Estimated Total Size (MB): 5.31\n",
       "============================================================================================================================================"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from monai.networks.nets import UNet\n",
    "from torchinfo import summary\n",
    "\n",
    "net = UNet(\n",
    "        spatial_dims=1,\n",
    "        in_channels=2,\n",
    "        out_channels=1,\n",
    "        channels=(16, 32, 64, 128, 256),\n",
    "        strides=(2, 2, 2, 2),\n",
    "        num_res_units=4\n",
    "    )\n",
    "\n",
    "summary(net,(1,2,224)) # bs, in_channel, vecter length"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "unet++"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BasicUNetPlusPlus features: (16, 32, 64, 128, 256, 16).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "====================================================================================================\n",
       "Layer (type:depth-idx)                             Output Shape              Param #\n",
       "====================================================================================================\n",
       "BasicUNetPlusPlus                                  [1, 1, 224]               --\n",
       "├─TwoConv: 1-1                                     [1, 16, 224]              --\n",
       "│    └─Convolution: 2-1                            [1, 16, 224]              --\n",
       "│    │    └─Conv1d: 3-1                            [1, 16, 224]              112\n",
       "│    │    └─ADN: 3-2                               [1, 16, 224]              32\n",
       "│    └─Convolution: 2-2                            [1, 16, 224]              --\n",
       "│    │    └─Conv1d: 3-3                            [1, 16, 224]              784\n",
       "│    │    └─ADN: 3-4                               [1, 16, 224]              32\n",
       "├─Down: 1-2                                        [1, 32, 112]              --\n",
       "│    └─MaxPool1d: 2-3                              [1, 16, 112]              --\n",
       "│    └─TwoConv: 2-4                                [1, 32, 112]              --\n",
       "│    │    └─Convolution: 3-5                       [1, 32, 112]              1,632\n",
       "│    │    └─Convolution: 3-6                       [1, 32, 112]              3,168\n",
       "├─UpCat: 1-3                                       [1, 16, 224]              --\n",
       "│    └─UpSample: 2-5                               [1, 32, 224]              --\n",
       "│    │    └─ConvTranspose1d: 3-7                   [1, 32, 224]              2,080\n",
       "│    └─TwoConv: 2-6                                [1, 16, 224]              --\n",
       "│    │    └─Convolution: 3-8                       [1, 16, 224]              2,352\n",
       "│    │    └─Convolution: 3-9                       [1, 16, 224]              816\n",
       "├─Down: 1-4                                        [1, 64, 56]               --\n",
       "│    └─MaxPool1d: 2-7                              [1, 32, 56]               --\n",
       "│    └─TwoConv: 2-8                                [1, 64, 56]               --\n",
       "│    │    └─Convolution: 3-10                      [1, 64, 56]               6,336\n",
       "│    │    └─Convolution: 3-11                      [1, 64, 56]               12,480\n",
       "├─UpCat: 1-5                                       [1, 32, 112]              --\n",
       "│    └─UpSample: 2-9                               [1, 32, 112]              --\n",
       "│    │    └─ConvTranspose1d: 3-12                  [1, 32, 112]              4,128\n",
       "│    └─TwoConv: 2-10                               [1, 32, 112]              --\n",
       "│    │    └─Convolution: 3-13                      [1, 32, 112]              6,240\n",
       "│    │    └─Convolution: 3-14                      [1, 32, 112]              3,168\n",
       "├─UpCat: 1-6                                       [1, 16, 224]              --\n",
       "│    └─UpSample: 2-11                              [1, 32, 224]              --\n",
       "│    │    └─ConvTranspose1d: 3-15                  [1, 32, 224]              2,080\n",
       "│    └─TwoConv: 2-12                               [1, 16, 224]              --\n",
       "│    │    └─Convolution: 3-16                      [1, 16, 224]              3,120\n",
       "│    │    └─Convolution: 3-17                      [1, 16, 224]              816\n",
       "├─Down: 1-7                                        [1, 128, 28]              --\n",
       "│    └─MaxPool1d: 2-13                             [1, 64, 28]               --\n",
       "│    └─TwoConv: 2-14                               [1, 128, 28]              --\n",
       "│    │    └─Convolution: 3-18                      [1, 128, 28]              24,960\n",
       "│    │    └─Convolution: 3-19                      [1, 128, 28]              49,536\n",
       "├─UpCat: 1-8                                       [1, 64, 56]               --\n",
       "│    └─UpSample: 2-15                              [1, 64, 56]               --\n",
       "│    │    └─ConvTranspose1d: 3-20                  [1, 64, 56]               16,448\n",
       "│    └─TwoConv: 2-16                               [1, 64, 56]               --\n",
       "│    │    └─Convolution: 3-21                      [1, 64, 56]               24,768\n",
       "│    │    └─Convolution: 3-22                      [1, 64, 56]               12,480\n",
       "├─UpCat: 1-9                                       [1, 32, 112]              --\n",
       "│    └─UpSample: 2-17                              [1, 32, 112]              --\n",
       "│    │    └─ConvTranspose1d: 3-23                  [1, 32, 112]              4,128\n",
       "│    └─TwoConv: 2-18                               [1, 32, 112]              --\n",
       "│    │    └─Convolution: 3-24                      [1, 32, 112]              9,312\n",
       "│    │    └─Convolution: 3-25                      [1, 32, 112]              3,168\n",
       "├─UpCat: 1-10                                      [1, 16, 224]              --\n",
       "│    └─UpSample: 2-19                              [1, 32, 224]              --\n",
       "│    │    └─ConvTranspose1d: 3-26                  [1, 32, 224]              2,080\n",
       "│    └─TwoConv: 2-20                               [1, 16, 224]              --\n",
       "│    │    └─Convolution: 3-27                      [1, 16, 224]              3,888\n",
       "│    │    └─Convolution: 3-28                      [1, 16, 224]              816\n",
       "├─Down: 1-11                                       [1, 256, 14]              --\n",
       "│    └─MaxPool1d: 2-21                             [1, 128, 14]              --\n",
       "│    └─TwoConv: 2-22                               [1, 256, 14]              --\n",
       "│    │    └─Convolution: 3-29                      [1, 256, 14]              99,072\n",
       "│    │    └─Convolution: 3-30                      [1, 256, 14]              197,376\n",
       "├─UpCat: 1-12                                      [1, 128, 28]              --\n",
       "│    └─UpSample: 2-23                              [1, 128, 28]              --\n",
       "│    │    └─ConvTranspose1d: 3-31                  [1, 128, 28]              65,664\n",
       "│    └─TwoConv: 2-24                               [1, 128, 28]              --\n",
       "│    │    └─Convolution: 3-32                      [1, 128, 28]              98,688\n",
       "│    │    └─Convolution: 3-33                      [1, 128, 28]              49,536\n",
       "├─UpCat: 1-13                                      [1, 64, 56]               --\n",
       "│    └─UpSample: 2-25                              [1, 64, 56]               --\n",
       "│    │    └─ConvTranspose1d: 3-34                  [1, 64, 56]               16,448\n",
       "│    └─TwoConv: 2-26                               [1, 64, 56]               --\n",
       "│    │    └─Convolution: 3-35                      [1, 64, 56]               37,056\n",
       "│    │    └─Convolution: 3-36                      [1, 64, 56]               12,480\n",
       "├─UpCat: 1-14                                      [1, 32, 112]              --\n",
       "│    └─UpSample: 2-27                              [1, 32, 112]              --\n",
       "│    │    └─ConvTranspose1d: 3-37                  [1, 32, 112]              4,128\n",
       "│    └─TwoConv: 2-28                               [1, 32, 112]              --\n",
       "│    │    └─Convolution: 3-38                      [1, 32, 112]              12,384\n",
       "│    │    └─Convolution: 3-39                      [1, 32, 112]              3,168\n",
       "├─UpCat: 1-15                                      [1, 16, 224]              --\n",
       "│    └─UpSample: 2-29                              [1, 32, 224]              --\n",
       "│    │    └─ConvTranspose1d: 3-40                  [1, 32, 224]              2,080\n",
       "│    └─TwoConv: 2-30                               [1, 16, 224]              --\n",
       "│    │    └─Convolution: 3-41                      [1, 16, 224]              4,656\n",
       "│    │    └─Convolution: 3-42                      [1, 16, 224]              816\n",
       "├─Conv1d: 1-16                                     [1, 1, 224]               17\n",
       "├─Conv1d: 1-17                                     [1, 1, 224]               17\n",
       "├─Conv1d: 1-18                                     [1, 1, 224]               17\n",
       "├─Conv1d: 1-19                                     [1, 1, 224]               17\n",
       "====================================================================================================\n",
       "Total params: 804,580\n",
       "Trainable params: 804,580\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 31.85\n",
       "====================================================================================================\n",
       "Input size (MB): 0.00\n",
       "Forward/backward pass size (MB): 2.13\n",
       "Params size (MB): 3.22\n",
       "Estimated Total Size (MB): 5.35\n",
       "===================================================================================================="
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from monai.networks.nets import BasicUnetPlusPlus\n",
    "from torchinfo import summary\n",
    "\n",
    "net = BasicUnetPlusPlus(\n",
    "    spatial_dims=1, \n",
    "    in_channels=2,\n",
    "    out_channels=1, \n",
    "    features=(16, 32, 64, 128, 256, 16), \n",
    "    deep_supervision=False, \n",
    "    act=('LeakyReLU', {'inplace': True, 'negative_slope': 0.1}), \n",
    "    norm=('instance', {'affine': True}), \n",
    "    bias=True, \n",
    "    dropout=0.0, \n",
    "    upsample='deconv'\n",
    "    )\n",
    "summary(net,(1,2,224)) # bs, in_channel, vecter length"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "attention unet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "=======================================================================================================================================\n",
       "Layer (type:depth-idx)                                                                Output Shape              Param #\n",
       "=======================================================================================================================================\n",
       "AttentionUnet                                                                         [1, 1, 224]               --\n",
       "├─Sequential: 1-1                                                                     [1, 1, 224]               --\n",
       "│    └─ConvBlock: 2-1                                                                 [1, 16, 224]              --\n",
       "│    │    └─Sequential: 3-1                                                           [1, 16, 224]              960\n",
       "│    └─AttentionLayer: 2-2                                                            [1, 16, 224]              --\n",
       "│    │    └─Sequential: 3-2                                                           [1, 32, 112]              675,804\n",
       "│    │    └─UpConv: 3-3                                                               [1, 16, 224]              1,584\n",
       "│    │    └─AttentionBlock: 3-4                                                       [1, 16, 224]              315\n",
       "│    │    └─Convolution: 3-5                                                          [1, 16, 224]              1,553\n",
       "│    └─Convolution: 2-3                                                               [1, 1, 224]               --\n",
       "│    │    └─Conv1d: 3-6                                                               [1, 1, 224]               17\n",
       "=======================================================================================================================================\n",
       "Total params: 680,233\n",
       "Trainable params: 680,233\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 19.21\n",
       "=======================================================================================================================================\n",
       "Input size (MB): 0.00\n",
       "Forward/backward pass size (MB): 1.27\n",
       "Params size (MB): 2.72\n",
       "Estimated Total Size (MB): 3.99\n",
       "======================================================================================================================================="
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from monai.networks.nets import AttentionUnet\n",
    "from torchinfo import summary\n",
    "\n",
    "net = AttentionUnet(\n",
    "    spatial_dims=1, \n",
    "    in_channels=2, \n",
    "    out_channels=1, \n",
    "    channels=(16, 32, 64, 128, 256),\n",
    "    strides=(2, 2, 2, 2),\n",
    "    kernel_size=3, \n",
    "    up_kernel_size=3, \n",
    "    dropout=0.0\n",
    "    )\n",
    "\n",
    "summary(net,(1,2,224)) # bs, in_channel, vecter length"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "swin unet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "=========================================================================================================\n",
       "Layer (type:depth-idx)                                  Output Shape              Param #\n",
       "=========================================================================================================\n",
       "SwinUNETR                                               [1, 1, 256, 256]          --\n",
       "├─SwinTransformer: 1-1                                  [1, 24, 128, 128]         --\n",
       "│    └─PatchEmbed: 2-1                                  [1, 24, 128, 128]         --\n",
       "│    │    └─Conv2d: 3-1                                 [1, 24, 128, 128]         216\n",
       "│    └─Dropout: 2-2                                     [1, 24, 128, 128]         --\n",
       "│    └─ModuleList: 2-3                                  --                        --\n",
       "│    │    └─BasicLayer: 3-2                             [1, 48, 64, 64]           20,262\n",
       "│    └─ModuleList: 2-4                                  --                        --\n",
       "│    │    └─BasicLayer: 3-3                             [1, 96, 32, 32]           77,388\n",
       "│    └─ModuleList: 2-5                                  --                        --\n",
       "│    │    └─BasicLayer: 3-4                             [1, 192, 16, 16]          302,232\n",
       "│    └─ModuleList: 2-6                                  --                        --\n",
       "│    │    └─BasicLayer: 3-5                             [1, 384, 8, 8]            1,194,288\n",
       "├─UnetrBasicBlock: 1-2                                  [1, 24, 256, 256]         --\n",
       "│    └─UnetResBlock: 2-7                                [1, 24, 256, 256]         --\n",
       "│    │    └─Convolution: 3-6                            [1, 24, 256, 256]         432\n",
       "│    │    └─InstanceNorm2d: 3-7                         [1, 24, 256, 256]         --\n",
       "│    │    └─LeakyReLU: 3-8                              [1, 24, 256, 256]         --\n",
       "│    │    └─Convolution: 3-9                            [1, 24, 256, 256]         5,184\n",
       "│    │    └─InstanceNorm2d: 3-10                        [1, 24, 256, 256]         --\n",
       "│    │    └─Convolution: 3-11                           [1, 24, 256, 256]         48\n",
       "│    │    └─InstanceNorm2d: 3-12                        [1, 24, 256, 256]         --\n",
       "│    │    └─LeakyReLU: 3-13                             [1, 24, 256, 256]         --\n",
       "├─UnetrBasicBlock: 1-3                                  [1, 24, 128, 128]         --\n",
       "│    └─UnetResBlock: 2-8                                [1, 24, 128, 128]         --\n",
       "│    │    └─Convolution: 3-14                           [1, 24, 128, 128]         5,184\n",
       "│    │    └─InstanceNorm2d: 3-15                        [1, 24, 128, 128]         --\n",
       "│    │    └─LeakyReLU: 3-16                             [1, 24, 128, 128]         --\n",
       "│    │    └─Convolution: 3-17                           [1, 24, 128, 128]         5,184\n",
       "│    │    └─InstanceNorm2d: 3-18                        [1, 24, 128, 128]         --\n",
       "│    │    └─LeakyReLU: 3-19                             [1, 24, 128, 128]         --\n",
       "├─UnetrBasicBlock: 1-4                                  [1, 48, 64, 64]           --\n",
       "│    └─UnetResBlock: 2-9                                [1, 48, 64, 64]           --\n",
       "│    │    └─Convolution: 3-20                           [1, 48, 64, 64]           20,736\n",
       "│    │    └─InstanceNorm2d: 3-21                        [1, 48, 64, 64]           --\n",
       "│    │    └─LeakyReLU: 3-22                             [1, 48, 64, 64]           --\n",
       "│    │    └─Convolution: 3-23                           [1, 48, 64, 64]           20,736\n",
       "│    │    └─InstanceNorm2d: 3-24                        [1, 48, 64, 64]           --\n",
       "│    │    └─LeakyReLU: 3-25                             [1, 48, 64, 64]           --\n",
       "├─UnetrBasicBlock: 1-5                                  [1, 96, 32, 32]           --\n",
       "│    └─UnetResBlock: 2-10                               [1, 96, 32, 32]           --\n",
       "│    │    └─Convolution: 3-26                           [1, 96, 32, 32]           82,944\n",
       "│    │    └─InstanceNorm2d: 3-27                        [1, 96, 32, 32]           --\n",
       "│    │    └─LeakyReLU: 3-28                             [1, 96, 32, 32]           --\n",
       "│    │    └─Convolution: 3-29                           [1, 96, 32, 32]           82,944\n",
       "│    │    └─InstanceNorm2d: 3-30                        [1, 96, 32, 32]           --\n",
       "│    │    └─LeakyReLU: 3-31                             [1, 96, 32, 32]           --\n",
       "├─UnetrBasicBlock: 1-6                                  [1, 384, 8, 8]            --\n",
       "│    └─UnetResBlock: 2-11                               [1, 384, 8, 8]            --\n",
       "│    │    └─Convolution: 3-32                           [1, 384, 8, 8]            1,327,104\n",
       "│    │    └─InstanceNorm2d: 3-33                        [1, 384, 8, 8]            --\n",
       "│    │    └─LeakyReLU: 3-34                             [1, 384, 8, 8]            --\n",
       "│    │    └─Convolution: 3-35                           [1, 384, 8, 8]            1,327,104\n",
       "│    │    └─InstanceNorm2d: 3-36                        [1, 384, 8, 8]            --\n",
       "│    │    └─LeakyReLU: 3-37                             [1, 384, 8, 8]            --\n",
       "├─UnetrUpBlock: 1-7                                     [1, 192, 16, 16]          --\n",
       "│    └─Convolution: 2-12                                [1, 192, 16, 16]          --\n",
       "│    │    └─ConvTranspose2d: 3-38                       [1, 192, 16, 16]          294,912\n",
       "│    └─UnetResBlock: 2-13                               [1, 192, 16, 16]          --\n",
       "│    │    └─Convolution: 3-39                           [1, 192, 16, 16]          663,552\n",
       "│    │    └─InstanceNorm2d: 3-40                        [1, 192, 16, 16]          --\n",
       "│    │    └─LeakyReLU: 3-41                             [1, 192, 16, 16]          --\n",
       "│    │    └─Convolution: 3-42                           [1, 192, 16, 16]          331,776\n",
       "│    │    └─InstanceNorm2d: 3-43                        [1, 192, 16, 16]          --\n",
       "│    │    └─Convolution: 3-44                           [1, 192, 16, 16]          73,728\n",
       "│    │    └─InstanceNorm2d: 3-45                        [1, 192, 16, 16]          --\n",
       "│    │    └─LeakyReLU: 3-46                             [1, 192, 16, 16]          --\n",
       "├─UnetrUpBlock: 1-8                                     [1, 96, 32, 32]           --\n",
       "│    └─Convolution: 2-14                                [1, 96, 32, 32]           --\n",
       "│    │    └─ConvTranspose2d: 3-47                       [1, 96, 32, 32]           73,728\n",
       "│    └─UnetResBlock: 2-15                               [1, 96, 32, 32]           --\n",
       "│    │    └─Convolution: 3-48                           [1, 96, 32, 32]           165,888\n",
       "│    │    └─InstanceNorm2d: 3-49                        [1, 96, 32, 32]           --\n",
       "│    │    └─LeakyReLU: 3-50                             [1, 96, 32, 32]           --\n",
       "│    │    └─Convolution: 3-51                           [1, 96, 32, 32]           82,944\n",
       "│    │    └─InstanceNorm2d: 3-52                        [1, 96, 32, 32]           --\n",
       "│    │    └─Convolution: 3-53                           [1, 96, 32, 32]           18,432\n",
       "│    │    └─InstanceNorm2d: 3-54                        [1, 96, 32, 32]           --\n",
       "│    │    └─LeakyReLU: 3-55                             [1, 96, 32, 32]           --\n",
       "├─UnetrUpBlock: 1-9                                     [1, 48, 64, 64]           --\n",
       "│    └─Convolution: 2-16                                [1, 48, 64, 64]           --\n",
       "│    │    └─ConvTranspose2d: 3-56                       [1, 48, 64, 64]           18,432\n",
       "│    └─UnetResBlock: 2-17                               [1, 48, 64, 64]           --\n",
       "│    │    └─Convolution: 3-57                           [1, 48, 64, 64]           41,472\n",
       "│    │    └─InstanceNorm2d: 3-58                        [1, 48, 64, 64]           --\n",
       "│    │    └─LeakyReLU: 3-59                             [1, 48, 64, 64]           --\n",
       "│    │    └─Convolution: 3-60                           [1, 48, 64, 64]           20,736\n",
       "│    │    └─InstanceNorm2d: 3-61                        [1, 48, 64, 64]           --\n",
       "│    │    └─Convolution: 3-62                           [1, 48, 64, 64]           4,608\n",
       "│    │    └─InstanceNorm2d: 3-63                        [1, 48, 64, 64]           --\n",
       "│    │    └─LeakyReLU: 3-64                             [1, 48, 64, 64]           --\n",
       "├─UnetrUpBlock: 1-10                                    [1, 24, 128, 128]         --\n",
       "│    └─Convolution: 2-18                                [1, 24, 128, 128]         --\n",
       "│    │    └─ConvTranspose2d: 3-65                       [1, 24, 128, 128]         4,608\n",
       "│    └─UnetResBlock: 2-19                               [1, 24, 128, 128]         --\n",
       "│    │    └─Convolution: 3-66                           [1, 24, 128, 128]         10,368\n",
       "│    │    └─InstanceNorm2d: 3-67                        [1, 24, 128, 128]         --\n",
       "│    │    └─LeakyReLU: 3-68                             [1, 24, 128, 128]         --\n",
       "│    │    └─Convolution: 3-69                           [1, 24, 128, 128]         5,184\n",
       "│    │    └─InstanceNorm2d: 3-70                        [1, 24, 128, 128]         --\n",
       "│    │    └─Convolution: 3-71                           [1, 24, 128, 128]         1,152\n",
       "│    │    └─InstanceNorm2d: 3-72                        [1, 24, 128, 128]         --\n",
       "│    │    └─LeakyReLU: 3-73                             [1, 24, 128, 128]         --\n",
       "├─UnetrUpBlock: 1-11                                    [1, 24, 256, 256]         --\n",
       "│    └─Convolution: 2-20                                [1, 24, 256, 256]         --\n",
       "│    │    └─ConvTranspose2d: 3-74                       [1, 24, 256, 256]         2,304\n",
       "│    └─UnetResBlock: 2-21                               [1, 24, 256, 256]         --\n",
       "│    │    └─Convolution: 3-75                           [1, 24, 256, 256]         10,368\n",
       "│    │    └─InstanceNorm2d: 3-76                        [1, 24, 256, 256]         --\n",
       "│    │    └─LeakyReLU: 3-77                             [1, 24, 256, 256]         --\n",
       "│    │    └─Convolution: 3-78                           [1, 24, 256, 256]         5,184\n",
       "│    │    └─InstanceNorm2d: 3-79                        [1, 24, 256, 256]         --\n",
       "│    │    └─Convolution: 3-80                           [1, 24, 256, 256]         1,152\n",
       "│    │    └─InstanceNorm2d: 3-81                        [1, 24, 256, 256]         --\n",
       "│    │    └─LeakyReLU: 3-82                             [1, 24, 256, 256]         --\n",
       "├─UnetOutBlock: 1-12                                    [1, 1, 256, 256]          --\n",
       "│    └─Convolution: 2-22                                [1, 1, 256, 256]          --\n",
       "│    │    └─Conv2d: 3-83                                [1, 1, 256, 256]          25\n",
       "=========================================================================================================\n",
       "Total params: 6,302,539\n",
       "Trainable params: 6,302,539\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.GIGABYTES): 3.71\n",
       "=========================================================================================================\n",
       "Input size (MB): 0.52\n",
       "Forward/backward pass size (MB): 273.34\n",
       "Params size (MB): 25.15\n",
       "Estimated Total Size (MB): 299.01\n",
       "========================================================================================================="
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from monai.networks.nets import SwinUNETR\n",
    "from torchinfo import summary\n",
    "\n",
    "net = SwinUNETR(\n",
    "    img_size=(256,256),\n",
    "    in_channels=2, \n",
    "    out_channels=1,\n",
    "    depths=(2, 2, 2, 2), \n",
    "    num_heads=(3, 6, 12, 24), \n",
    "    feature_size=24, \n",
    "    norm_name='instance', \n",
    "    drop_rate=0.0, \n",
    "    attn_drop_rate=0.0, \n",
    "    dropout_path_rate=0.0, \n",
    "    normalize=True, \n",
    "    use_checkpoint=False, \n",
    "    spatial_dims=2, \n",
    "    downsample='merging', \n",
    "    use_v2=False\n",
    "    )\n",
    "summary(net,(1,2,256,256)) # bs, in_channel, vecter length"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "transformer unet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "====================================================================================================\n",
       "Layer (type:depth-idx)                             Output Shape              Param #\n",
       "====================================================================================================\n",
       "UNETR                                              [1, 1, 256, 256]          --\n",
       "├─ViT: 1-1                                         [1, 256, 768]             --\n",
       "│    └─PatchEmbeddingBlock: 2-1                    [1, 256, 768]             196,608\n",
       "│    │    └─Conv2d: 3-1                            [1, 768, 16, 16]          590,592\n",
       "│    │    └─Dropout: 3-2                           [1, 256, 768]             --\n",
       "│    └─ModuleList: 2-2                             --                        --\n",
       "│    │    └─TransformerBlock: 3-3                  [1, 256, 768]             7,085,568\n",
       "│    │    └─TransformerBlock: 3-4                  [1, 256, 768]             7,085,568\n",
       "│    │    └─TransformerBlock: 3-5                  [1, 256, 768]             7,085,568\n",
       "│    │    └─TransformerBlock: 3-6                  [1, 256, 768]             7,085,568\n",
       "│    │    └─TransformerBlock: 3-7                  [1, 256, 768]             7,085,568\n",
       "│    │    └─TransformerBlock: 3-8                  [1, 256, 768]             7,085,568\n",
       "│    │    └─TransformerBlock: 3-9                  [1, 256, 768]             7,085,568\n",
       "│    │    └─TransformerBlock: 3-10                 [1, 256, 768]             7,085,568\n",
       "│    │    └─TransformerBlock: 3-11                 [1, 256, 768]             7,085,568\n",
       "│    │    └─TransformerBlock: 3-12                 [1, 256, 768]             7,085,568\n",
       "│    │    └─TransformerBlock: 3-13                 [1, 256, 768]             7,085,568\n",
       "│    │    └─TransformerBlock: 3-14                 [1, 256, 768]             7,085,568\n",
       "│    └─LayerNorm: 2-3                              [1, 256, 768]             1,536\n",
       "├─UnetrBasicBlock: 1-2                             [1, 16, 256, 256]         --\n",
       "│    └─UnetResBlock: 2-4                           [1, 16, 256, 256]         --\n",
       "│    │    └─Convolution: 3-15                      [1, 16, 256, 256]         432\n",
       "│    │    └─InstanceNorm2d: 3-16                   [1, 16, 256, 256]         --\n",
       "│    │    └─LeakyReLU: 3-17                        [1, 16, 256, 256]         --\n",
       "│    │    └─Convolution: 3-18                      [1, 16, 256, 256]         2,304\n",
       "│    │    └─InstanceNorm2d: 3-19                   [1, 16, 256, 256]         --\n",
       "│    │    └─Convolution: 3-20                      [1, 16, 256, 256]         48\n",
       "│    │    └─InstanceNorm2d: 3-21                   [1, 16, 256, 256]         --\n",
       "│    │    └─LeakyReLU: 3-22                        [1, 16, 256, 256]         --\n",
       "├─UnetrPrUpBlock: 1-3                              [1, 32, 128, 128]         --\n",
       "│    └─Convolution: 2-5                            [1, 32, 32, 32]           --\n",
       "│    │    └─ConvTranspose2d: 3-23                  [1, 32, 32, 32]           98,304\n",
       "│    └─ModuleList: 2-6                             --                        --\n",
       "│    │    └─Sequential: 3-24                       [1, 32, 64, 64]           22,528\n",
       "│    │    └─Sequential: 3-25                       [1, 32, 128, 128]         22,528\n",
       "├─UnetrPrUpBlock: 1-4                              [1, 64, 64, 64]           --\n",
       "│    └─Convolution: 2-7                            [1, 64, 32, 32]           --\n",
       "│    │    └─ConvTranspose2d: 3-26                  [1, 64, 32, 32]           196,608\n",
       "│    └─ModuleList: 2-8                             --                        --\n",
       "│    │    └─Sequential: 3-27                       [1, 64, 64, 64]           90,112\n",
       "├─UnetrPrUpBlock: 1-5                              [1, 128, 32, 32]          --\n",
       "│    └─Convolution: 2-9                            [1, 128, 32, 32]          --\n",
       "│    │    └─ConvTranspose2d: 3-28                  [1, 128, 32, 32]          393,216\n",
       "├─UnetrUpBlock: 1-6                                [1, 128, 32, 32]          --\n",
       "│    └─Convolution: 2-10                           [1, 128, 32, 32]          --\n",
       "│    │    └─ConvTranspose2d: 3-29                  [1, 128, 32, 32]          393,216\n",
       "│    └─UnetResBlock: 2-11                          [1, 128, 32, 32]          --\n",
       "│    │    └─Convolution: 3-30                      [1, 128, 32, 32]          294,912\n",
       "│    │    └─InstanceNorm2d: 3-31                   [1, 128, 32, 32]          --\n",
       "│    │    └─LeakyReLU: 3-32                        [1, 128, 32, 32]          --\n",
       "│    │    └─Convolution: 3-33                      [1, 128, 32, 32]          147,456\n",
       "│    │    └─InstanceNorm2d: 3-34                   [1, 128, 32, 32]          --\n",
       "│    │    └─Convolution: 3-35                      [1, 128, 32, 32]          32,768\n",
       "│    │    └─InstanceNorm2d: 3-36                   [1, 128, 32, 32]          --\n",
       "│    │    └─LeakyReLU: 3-37                        [1, 128, 32, 32]          --\n",
       "├─UnetrUpBlock: 1-7                                [1, 64, 64, 64]           --\n",
       "│    └─Convolution: 2-12                           [1, 64, 64, 64]           --\n",
       "│    │    └─ConvTranspose2d: 3-38                  [1, 64, 64, 64]           32,768\n",
       "│    └─UnetResBlock: 2-13                          [1, 64, 64, 64]           --\n",
       "│    │    └─Convolution: 3-39                      [1, 64, 64, 64]           73,728\n",
       "│    │    └─InstanceNorm2d: 3-40                   [1, 64, 64, 64]           --\n",
       "│    │    └─LeakyReLU: 3-41                        [1, 64, 64, 64]           --\n",
       "│    │    └─Convolution: 3-42                      [1, 64, 64, 64]           36,864\n",
       "│    │    └─InstanceNorm2d: 3-43                   [1, 64, 64, 64]           --\n",
       "│    │    └─Convolution: 3-44                      [1, 64, 64, 64]           8,192\n",
       "│    │    └─InstanceNorm2d: 3-45                   [1, 64, 64, 64]           --\n",
       "│    │    └─LeakyReLU: 3-46                        [1, 64, 64, 64]           --\n",
       "├─UnetrUpBlock: 1-8                                [1, 32, 128, 128]         --\n",
       "│    └─Convolution: 2-14                           [1, 32, 128, 128]         --\n",
       "│    │    └─ConvTranspose2d: 3-47                  [1, 32, 128, 128]         8,192\n",
       "│    └─UnetResBlock: 2-15                          [1, 32, 128, 128]         --\n",
       "│    │    └─Convolution: 3-48                      [1, 32, 128, 128]         18,432\n",
       "│    │    └─InstanceNorm2d: 3-49                   [1, 32, 128, 128]         --\n",
       "│    │    └─LeakyReLU: 3-50                        [1, 32, 128, 128]         --\n",
       "│    │    └─Convolution: 3-51                      [1, 32, 128, 128]         9,216\n",
       "│    │    └─InstanceNorm2d: 3-52                   [1, 32, 128, 128]         --\n",
       "│    │    └─Convolution: 3-53                      [1, 32, 128, 128]         2,048\n",
       "│    │    └─InstanceNorm2d: 3-54                   [1, 32, 128, 128]         --\n",
       "│    │    └─LeakyReLU: 3-55                        [1, 32, 128, 128]         --\n",
       "├─UnetrUpBlock: 1-9                                [1, 16, 256, 256]         --\n",
       "│    └─Convolution: 2-16                           [1, 16, 256, 256]         --\n",
       "│    │    └─ConvTranspose2d: 3-56                  [1, 16, 256, 256]         2,048\n",
       "│    └─UnetResBlock: 2-17                          [1, 16, 256, 256]         --\n",
       "│    │    └─Convolution: 3-57                      [1, 16, 256, 256]         4,608\n",
       "│    │    └─InstanceNorm2d: 3-58                   [1, 16, 256, 256]         --\n",
       "│    │    └─LeakyReLU: 3-59                        [1, 16, 256, 256]         --\n",
       "│    │    └─Convolution: 3-60                      [1, 16, 256, 256]         2,304\n",
       "│    │    └─InstanceNorm2d: 3-61                   [1, 16, 256, 256]         --\n",
       "│    │    └─Convolution: 3-62                      [1, 16, 256, 256]         512\n",
       "│    │    └─InstanceNorm2d: 3-63                   [1, 16, 256, 256]         --\n",
       "│    │    └─LeakyReLU: 3-64                        [1, 16, 256, 256]         --\n",
       "├─UnetOutBlock: 1-10                               [1, 1, 256, 256]          --\n",
       "│    └─Convolution: 2-18                           [1, 1, 256, 256]          --\n",
       "│    │    └─Conv2d: 3-65                           [1, 1, 256, 256]          17\n",
       "====================================================================================================\n",
       "Total params: 87,708,913\n",
       "Trainable params: 87,708,913\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.GIGABYTES): 4.71\n",
       "====================================================================================================\n",
       "Input size (MB): 0.79\n",
       "Forward/backward pass size (MB): 323.22\n",
       "Params size (MB): 350.05\n",
       "Estimated Total Size (MB): 674.06\n",
       "===================================================================================================="
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from monai.networks.nets import UNETR\n",
    "from torchinfo import summary\n",
    "\n",
    "net = UNETR(\n",
    "    in_channels=3,\n",
    "    out_channels=1, \n",
    "    img_size=(256,256),\n",
    "    feature_size=16, \n",
    "    hidden_size=768, \n",
    "    mlp_dim=3072, \n",
    "    num_heads=12, \n",
    "    proj_type='conv', \n",
    "    norm_name='instance', \n",
    "    conv_block=True, \n",
    "    res_block=True, \n",
    "    dropout_rate=0.0, \n",
    "    spatial_dims=2, \n",
    "    qkv_bias=False, \n",
    "    save_attn=False)\n",
    "\n",
    "summary(net,(1,3,256,256))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "DiffusionModelUNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "timesteps:tensor([4])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "=========================================================================================================\n",
       "Layer (type:depth-idx)                                  Output Shape              Param #\n",
       "=========================================================================================================\n",
       "DiffusionModelUNet                                      [1, 1, 128, 128, 128]     --\n",
       "├─Sequential: 1-1                                       [1, 128]                  --\n",
       "│    └─Linear: 2-1                                      [1, 128]                  4,224\n",
       "│    └─SiLU: 2-2                                        [1, 128]                  --\n",
       "│    └─Linear: 2-3                                      [1, 128]                  16,512\n",
       "├─Convolution: 1-2                                      [1, 32, 128, 128, 128]    --\n",
       "│    └─Conv3d: 2-4                                      [1, 32, 128, 128, 128]    896\n",
       "├─ModuleList: 1-3                                       --                        --\n",
       "│    └─DownBlock: 2-5                                   [1, 32, 64, 64, 64]       --\n",
       "│    │    └─ModuleList: 3-1                             --                        119,232\n",
       "│    │    └─DiffusionUnetDownsample: 3-2                [1, 32, 64, 64, 64]       27,680\n",
       "│    └─DownBlock: 2-6                                   [1, 32, 32, 32, 32]       --\n",
       "│    │    └─ModuleList: 3-3                             --                        119,232\n",
       "│    │    └─DiffusionUnetDownsample: 3-4                [1, 32, 32, 32, 32]       27,680\n",
       "│    └─AttnDownBlock: 2-7                               [1, 32, 16, 16, 16]       --\n",
       "│    │    └─ModuleList: 3-7                             --                        (recursive)\n",
       "│    │    └─ModuleList: 3-8                             --                        (recursive)\n",
       "│    │    └─ModuleList: 3-7                             --                        (recursive)\n",
       "│    │    └─ModuleList: 3-8                             --                        (recursive)\n",
       "│    │    └─DiffusionUnetDownsample: 3-9                [1, 32, 16, 16, 16]       27,680\n",
       "│    └─AttnDownBlock: 2-8                               [1, 32, 16, 16, 16]       --\n",
       "│    │    └─ModuleList: 3-12                            --                        (recursive)\n",
       "│    │    └─ModuleList: 3-13                            --                        (recursive)\n",
       "│    │    └─ModuleList: 3-12                            --                        (recursive)\n",
       "│    │    └─ModuleList: 3-13                            --                        (recursive)\n",
       "├─AttnMidBlock: 1-4                                     [1, 32, 16, 16, 16]       --\n",
       "│    └─DiffusionUNetResnetBlock: 2-9                    [1, 32, 16, 16, 16]       --\n",
       "│    │    └─GroupNorm: 3-14                             [1, 32, 16, 16, 16]       64\n",
       "│    │    └─SiLU: 3-15                                  [1, 32, 16, 16, 16]       --\n",
       "│    │    └─Convolution: 3-16                           [1, 32, 16, 16, 16]       27,680\n",
       "│    │    └─SiLU: 3-17                                  [1, 128]                  --\n",
       "│    │    └─Linear: 3-18                                [1, 32]                   4,128\n",
       "│    │    └─GroupNorm: 3-19                             [1, 32, 16, 16, 16]       64\n",
       "│    │    └─SiLU: 3-20                                  [1, 32, 16, 16, 16]       --\n",
       "│    │    └─Convolution: 3-21                           [1, 32, 16, 16, 16]       27,680\n",
       "│    │    └─Identity: 3-22                              [1, 32, 16, 16, 16]       --\n",
       "│    └─SpatialAttentionBlock: 2-10                      [1, 32, 16, 16, 16]       --\n",
       "│    │    └─GroupNorm: 3-23                             [1, 32, 16, 16, 16]       64\n",
       "│    │    └─SABlock: 3-24                               [1, 4096, 32]             4,224\n",
       "│    └─DiffusionUNetResnetBlock: 2-11                   [1, 32, 16, 16, 16]       --\n",
       "│    │    └─GroupNorm: 3-25                             [1, 32, 16, 16, 16]       64\n",
       "│    │    └─SiLU: 3-26                                  [1, 32, 16, 16, 16]       --\n",
       "│    │    └─Convolution: 3-27                           [1, 32, 16, 16, 16]       27,680\n",
       "│    │    └─SiLU: 3-28                                  [1, 128]                  --\n",
       "│    │    └─Linear: 3-29                                [1, 32]                   4,128\n",
       "│    │    └─GroupNorm: 3-30                             [1, 32, 16, 16, 16]       64\n",
       "│    │    └─SiLU: 3-31                                  [1, 32, 16, 16, 16]       --\n",
       "│    │    └─Convolution: 3-32                           [1, 32, 16, 16, 16]       27,680\n",
       "│    │    └─Identity: 3-33                              [1, 32, 16, 16, 16]       --\n",
       "├─ModuleList: 1-5                                       --                        --\n",
       "│    └─AttnUpBlock: 2-12                                [1, 32, 32, 32, 32]       --\n",
       "│    │    └─ModuleList: 3-38                            --                        (recursive)\n",
       "│    │    └─ModuleList: 3-39                            --                        (recursive)\n",
       "│    │    └─ModuleList: 3-38                            --                        (recursive)\n",
       "│    │    └─ModuleList: 3-39                            --                        (recursive)\n",
       "│    │    └─ModuleList: 3-38                            --                        (recursive)\n",
       "│    │    └─ModuleList: 3-39                            --                        (recursive)\n",
       "│    │    └─WrappedUpsample: 3-40                       [1, 32, 32, 32, 32]       27,680\n",
       "│    └─AttnUpBlock: 2-13                                [1, 32, 64, 64, 64]       --\n",
       "│    │    └─ModuleList: 3-45                            --                        (recursive)\n",
       "│    │    └─ModuleList: 3-46                            --                        (recursive)\n",
       "│    │    └─ModuleList: 3-45                            --                        (recursive)\n",
       "│    │    └─ModuleList: 3-46                            --                        (recursive)\n",
       "│    │    └─ModuleList: 3-45                            --                        (recursive)\n",
       "│    │    └─ModuleList: 3-46                            --                        (recursive)\n",
       "│    │    └─WrappedUpsample: 3-47                       [1, 32, 64, 64, 64]       27,680\n",
       "│    └─UpBlock: 2-14                                    [1, 32, 128, 128, 128]    --\n",
       "│    │    └─ModuleList: 3-48                            --                        268,224\n",
       "│    │    └─WrappedUpsample: 3-49                       [1, 32, 128, 128, 128]    27,680\n",
       "│    └─UpBlock: 2-15                                    [1, 32, 128, 128, 128]    --\n",
       "│    │    └─ModuleList: 3-50                            --                        268,224\n",
       "├─Sequential: 1-6                                       [1, 1, 128, 128, 128]     --\n",
       "│    └─GroupNorm: 2-16                                  [1, 32, 128, 128, 128]    64\n",
       "│    └─SiLU: 2-17                                       [1, 32, 128, 128, 128]    --\n",
       "│    └─Convolution: 2-18                                [1, 1, 128, 128, 128]     --\n",
       "│    │    └─Conv3d: 3-51                                [1, 1, 128, 128, 128]     865\n",
       "=========================================================================================================\n",
       "Total params: 1,904,865\n",
       "Trainable params: 1,904,865\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.GIGABYTES): 955.59\n",
       "=========================================================================================================\n",
       "Input size (MB): 8.39\n",
       "Forward/backward pass size (MB): 17977.84\n",
       "Params size (MB): 7.62\n",
       "Estimated Total Size (MB): 17993.85\n",
       "========================================================================================================="
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from monai.networks.nets import DiffusionModelUNet\n",
    "from torchinfo import summary\n",
    "import torch\n",
    "\n",
    "net = DiffusionModelUNet(\n",
    "    spatial_dims=3, \n",
    "    in_channels=1, \n",
    "    out_channels=1, \n",
    "    channels=(32, 32, 32, 32), \n",
    "    transformer_num_layers= 1,\n",
    "    )\n",
    "\n",
    "x = torch.randn(1, 1, 128, 128, 128)  # bs, in_channel, vecter length\n",
    "timesteps = torch.randint(1, 5, (1,))  # timestep between 0-10\n",
    "print(f'timesteps:{timesteps}')\n",
    "summary(net, input_data=(x, timesteps))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SegResNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==========================================================================================\n",
       "Layer (type:depth-idx)                   Output Shape              Param #\n",
       "==========================================================================================\n",
       "SegResNet                                [1, 1, 256, 256]          --\n",
       "├─Convolution: 1-1                       [1, 8, 256, 256]          --\n",
       "│    └─Conv2d: 2-1                       [1, 8, 256, 256]          216\n",
       "├─ModuleList: 1-2                        --                        --\n",
       "│    └─Sequential: 2-2                   [1, 8, 256, 256]          --\n",
       "│    │    └─Identity: 3-1                [1, 8, 256, 256]          --\n",
       "│    │    └─ResBlock: 3-2                [1, 8, 256, 256]          1,184\n",
       "│    └─Sequential: 2-3                   [1, 16, 128, 128]         --\n",
       "│    │    └─Convolution: 3-3             [1, 16, 128, 128]         1,152\n",
       "│    │    └─ResBlock: 3-4                [1, 16, 128, 128]         4,672\n",
       "│    │    └─ResBlock: 3-5                [1, 16, 128, 128]         4,672\n",
       "│    └─Sequential: 2-4                   [1, 32, 64, 64]           --\n",
       "│    │    └─Convolution: 3-6             [1, 32, 64, 64]           4,608\n",
       "│    │    └─ResBlock: 3-7                [1, 32, 64, 64]           18,560\n",
       "│    │    └─ResBlock: 3-8                [1, 32, 64, 64]           18,560\n",
       "│    └─Sequential: 2-5                   [1, 64, 32, 32]           --\n",
       "│    │    └─Convolution: 3-9             [1, 64, 32, 32]           18,432\n",
       "│    │    └─ResBlock: 3-10               [1, 64, 32, 32]           73,984\n",
       "│    │    └─ResBlock: 3-11               [1, 64, 32, 32]           73,984\n",
       "│    │    └─ResBlock: 3-12               [1, 64, 32, 32]           73,984\n",
       "│    │    └─ResBlock: 3-13               [1, 64, 32, 32]           73,984\n",
       "├─ModuleList: 1-7                        --                        (recursive)\n",
       "│    └─Sequential: 2-6                   [1, 32, 64, 64]           --\n",
       "│    │    └─Convolution: 3-14            [1, 32, 32, 32]           2,048\n",
       "│    │    └─UpSample: 3-15               [1, 32, 64, 64]           --\n",
       "├─ModuleList: 1-8                        --                        (recursive)\n",
       "│    └─Sequential: 2-7                   [1, 32, 64, 64]           --\n",
       "│    │    └─ResBlock: 3-16               [1, 32, 64, 64]           18,560\n",
       "├─ModuleList: 1-7                        --                        (recursive)\n",
       "│    └─Sequential: 2-8                   [1, 16, 128, 128]         --\n",
       "│    │    └─Convolution: 3-17            [1, 16, 64, 64]           512\n",
       "│    │    └─UpSample: 3-18               [1, 16, 128, 128]         --\n",
       "├─ModuleList: 1-8                        --                        (recursive)\n",
       "│    └─Sequential: 2-9                   [1, 16, 128, 128]         --\n",
       "│    │    └─ResBlock: 3-19               [1, 16, 128, 128]         4,672\n",
       "├─ModuleList: 1-7                        --                        (recursive)\n",
       "│    └─Sequential: 2-10                  [1, 8, 256, 256]          --\n",
       "│    │    └─Convolution: 3-20            [1, 8, 128, 128]          128\n",
       "│    │    └─UpSample: 3-21               [1, 8, 256, 256]          --\n",
       "├─ModuleList: 1-8                        --                        (recursive)\n",
       "│    └─Sequential: 2-11                  [1, 8, 256, 256]          --\n",
       "│    │    └─ResBlock: 3-22               [1, 8, 256, 256]          1,184\n",
       "├─Sequential: 1-9                        [1, 1, 256, 256]          --\n",
       "│    └─GroupNorm: 2-12                   [1, 8, 256, 256]          16\n",
       "│    └─ReLU: 2-13                        [1, 8, 256, 256]          --\n",
       "│    └─Convolution: 2-14                 [1, 1, 256, 256]          --\n",
       "│    │    └─Conv2d: 3-23                 [1, 1, 256, 256]          9\n",
       "==========================================================================================\n",
       "Total params: 395,121\n",
       "Trainable params: 395,121\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 983.63\n",
       "==========================================================================================\n",
       "Input size (MB): 0.79\n",
       "Forward/backward pass size (MB): 94.11\n",
       "Params size (MB): 1.58\n",
       "Estimated Total Size (MB): 96.48\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from monai.networks.nets import SegResNet\n",
    "from torchinfo import summary\n",
    "\n",
    "\n",
    "net = SegResNet(\n",
    "    spatial_dims=2, \n",
    "    init_filters=8, \n",
    "    in_channels=3, \n",
    "    out_channels=1, \n",
    "    dropout_prob=None, \n",
    "    act=('RELU', {'inplace': True}), \n",
    "    norm=('GROUP', {'num_groups': 8}), \n",
    "    norm_name='', \n",
    "    num_groups=8, \n",
    "    use_conv_final=True, \n",
    "    blocks_down=(1, 2, 2, 4), \n",
    "    blocks_up=(1, 1, 1), \n",
    "    )\n",
    "\n",
    "summary(net,(1,3,256,256))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SegResNetVAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==========================================================================================\n",
       "Layer (type:depth-idx)                   Output Shape              Param #\n",
       "==========================================================================================\n",
       "SegResNetVAE                             [1, 1, 256, 256]          3,160,871\n",
       "├─Convolution: 1-1                       [1, 8, 256, 256]          --\n",
       "│    └─Conv2d: 2-1                       [1, 8, 256, 256]          216\n",
       "├─ModuleList: 1-2                        --                        --\n",
       "│    └─Sequential: 2-2                   [1, 8, 256, 256]          --\n",
       "│    │    └─Identity: 3-1                [1, 8, 256, 256]          --\n",
       "│    │    └─ResBlock: 3-2                [1, 8, 256, 256]          1,184\n",
       "│    └─Sequential: 2-3                   [1, 16, 128, 128]         --\n",
       "│    │    └─Convolution: 3-3             [1, 16, 128, 128]         1,152\n",
       "│    │    └─ResBlock: 3-4                [1, 16, 128, 128]         4,672\n",
       "│    │    └─ResBlock: 3-5                [1, 16, 128, 128]         4,672\n",
       "│    └─Sequential: 2-4                   [1, 32, 64, 64]           --\n",
       "│    │    └─Convolution: 3-6             [1, 32, 64, 64]           4,608\n",
       "│    │    └─ResBlock: 3-7                [1, 32, 64, 64]           18,560\n",
       "│    │    └─ResBlock: 3-8                [1, 32, 64, 64]           18,560\n",
       "│    └─Sequential: 2-5                   [1, 64, 32, 32]           --\n",
       "│    │    └─Convolution: 3-9             [1, 64, 32, 32]           18,432\n",
       "│    │    └─ResBlock: 3-10               [1, 64, 32, 32]           73,984\n",
       "│    │    └─ResBlock: 3-11               [1, 64, 32, 32]           73,984\n",
       "│    │    └─ResBlock: 3-12               [1, 64, 32, 32]           73,984\n",
       "│    │    └─ResBlock: 3-13               [1, 64, 32, 32]           73,984\n",
       "├─ModuleList: 1-7                        --                        (recursive)\n",
       "│    └─Sequential: 2-6                   [1, 32, 64, 64]           --\n",
       "│    │    └─Convolution: 3-14            [1, 32, 32, 32]           2,048\n",
       "│    │    └─UpSample: 3-15               [1, 32, 64, 64]           --\n",
       "├─ModuleList: 1-8                        --                        (recursive)\n",
       "│    └─Sequential: 2-7                   [1, 32, 64, 64]           --\n",
       "│    │    └─ResBlock: 3-16               [1, 32, 64, 64]           18,560\n",
       "├─ModuleList: 1-7                        --                        (recursive)\n",
       "│    └─Sequential: 2-8                   [1, 16, 128, 128]         --\n",
       "│    │    └─Convolution: 3-17            [1, 16, 64, 64]           512\n",
       "│    │    └─UpSample: 3-18               [1, 16, 128, 128]         --\n",
       "├─ModuleList: 1-8                        --                        (recursive)\n",
       "│    └─Sequential: 2-9                   [1, 16, 128, 128]         --\n",
       "│    │    └─ResBlock: 3-19               [1, 16, 128, 128]         4,672\n",
       "├─ModuleList: 1-7                        --                        (recursive)\n",
       "│    └─Sequential: 2-10                  [1, 8, 256, 256]          --\n",
       "│    │    └─Convolution: 3-20            [1, 8, 128, 128]          128\n",
       "│    │    └─UpSample: 3-21               [1, 8, 256, 256]          --\n",
       "├─ModuleList: 1-8                        --                        (recursive)\n",
       "│    └─Sequential: 2-11                  [1, 8, 256, 256]          --\n",
       "│    │    └─ResBlock: 3-22               [1, 8, 256, 256]          1,184\n",
       "├─Sequential: 1-9                        [1, 1, 256, 256]          9\n",
       "│    └─GroupNorm: 2-12                   [1, 8, 256, 256]          16\n",
       "├─Sequential: 1-10                       --                        43\n",
       "│    └─ReLU: 2-13                        [1, 8, 256, 256]          --\n",
       "├─Sequential: 1-11                       --                        (recursive)\n",
       "│    └─Convolution: 2-14                 [1, 1, 256, 256]          --\n",
       "│    │    └─Conv2d: 3-23                 [1, 1, 256, 256]          9\n",
       "==========================================================================================\n",
       "Total params: 3,556,044\n",
       "Trainable params: 3,556,044\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 983.63\n",
       "==========================================================================================\n",
       "Input size (MB): 0.79\n",
       "Forward/backward pass size (MB): 94.11\n",
       "Params size (MB): 1.58\n",
       "Estimated Total Size (MB): 96.48\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from monai.networks.nets import SegResNetVAE \n",
    "from torchinfo import summary\n",
    "\n",
    "net = SegResNetVAE(\n",
    "    input_image_size = (256, 256), \n",
    "    vae_estimate_std=False, \n",
    "    vae_default_std=0.3, \n",
    "    vae_nz=256, \n",
    "    spatial_dims=2, \n",
    "    init_filters=8, \n",
    "    in_channels=3,\n",
    "    out_channels=1, \n",
    "    dropout_prob=None, \n",
    "    act=('RELU', {'inplace': True}), \n",
    "    norm=('GROUP', {'num_groups': 8}), \n",
    "    use_conv_final=True, \n",
    "    blocks_down=(1, 2, 2, 4), \n",
    "    blocks_up=(1, 1, 1), \n",
    "    )\n",
    "\n",
    "summary(net,(1,3,256,256))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SegResNetDS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "====================================================================================================\n",
       "Layer (type:depth-idx)                             Output Shape              Param #\n",
       "====================================================================================================\n",
       "SegResNetDS                                        [1, 1, 256, 256]          --\n",
       "├─SegResEncoder: 1-1                               [1, 32, 256, 256]         --\n",
       "│    └─Conv2d: 2-1                                 [1, 32, 256, 256]         864\n",
       "│    └─ModuleList: 2-2                             --                        --\n",
       "│    │    └─ModuleDict: 3-1                        --                        36,992\n",
       "│    │    └─ModuleDict: 3-2                        --                        221,696\n",
       "│    │    └─ModuleDict: 3-3                        --                        885,760\n",
       "│    │    └─ModuleDict: 3-4                        --                        4,722,688\n",
       "├─ModuleList: 1-2                                  --                        --\n",
       "│    └─ModuleDict: 2-3                             --                        --\n",
       "│    │    └─UpSample: 3-5                          [1, 128, 64, 64]          294,912\n",
       "│    │    └─Sequential: 3-6                        [1, 128, 64, 64]          295,424\n",
       "│    └─ModuleDict: 2-4                             --                        --\n",
       "│    │    └─UpSample: 3-7                          [1, 64, 128, 128]         73,728\n",
       "│    │    └─Sequential: 3-8                        [1, 64, 128, 128]         73,984\n",
       "│    └─ModuleDict: 2-5                             --                        --\n",
       "│    │    └─UpSample: 3-9                          [1, 32, 256, 256]         18,432\n",
       "│    │    └─Sequential: 3-10                       [1, 32, 256, 256]         18,560\n",
       "│    │    └─Conv2d: 3-11                           [1, 1, 256, 256]          33\n",
       "====================================================================================================\n",
       "Total params: 6,643,073\n",
       "Trainable params: 6,643,073\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.GIGABYTES): 19.08\n",
       "====================================================================================================\n",
       "Input size (MB): 0.79\n",
       "Forward/backward pass size (MB): 380.11\n",
       "Params size (MB): 26.57\n",
       "Estimated Total Size (MB): 407.47\n",
       "===================================================================================================="
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from monai.networks.nets import SegResNetDS\n",
    "from torchinfo import summary\n",
    "\n",
    "\n",
    "net = SegResNetDS(\n",
    "    spatial_dims=2, \n",
    "    init_filters=32, \n",
    "    in_channels=3, \n",
    "    out_channels=1, \n",
    "    act='relu', \n",
    "    norm='batch', \n",
    "    blocks_down=(1, 2, 2, 4), \n",
    "    blocks_up=None, \n",
    "    dsdepth=1, \n",
    "    preprocess=None, \n",
    "    upsample_mode='deconv',\n",
    "    resolution=None\n",
    "    )\n",
    "\n",
    "summary(net,(1,3,256,256))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## models from Segmentation Models Pytroch 3D\n",
    "* Reference: https://github.com/ZFTurbo/segmentation_models_pytorch_3d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "functools.partial(<function preprocess_input at 0x000001BD6982F740>, input_space='RGB', input_range=[0, 1], mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "=========================================================================================================\n",
       "Layer (type:depth-idx)                                  Output Shape              Param #\n",
       "=========================================================================================================\n",
       "Unet                                                    [1, 3, 256, 256, 256]     --\n",
       "├─EfficientNetEncoder: 1-1                              [1, 3, 256, 256, 256]     412,160\n",
       "│    └─Conv3dStaticSamePadding: 2-1                     [1, 32, 128, 128, 128]    2,592\n",
       "│    │    └─ZeroPad2d: 3-1                              [1, 3, 257, 257, 257]     --\n",
       "│    └─BatchNorm3d: 2-2                                 [1, 32, 128, 128, 128]    64\n",
       "│    └─MemoryEfficientSwish: 2-3                        [1, 32, 128, 128, 128]    --\n",
       "│    └─ModuleList: 2-4                                  --                        --\n",
       "│    │    └─MBConvBlock: 3-2                            [1, 16, 128, 128, 128]    2,024\n",
       "│    │    └─MBConvBlock: 3-3                            [1, 24, 64, 64, 64]       7,732\n",
       "│    │    └─MBConvBlock: 3-4                            [1, 24, 64, 64, 64]       13,302\n",
       "│    │    └─MBConvBlock: 3-5                            [1, 40, 32, 32, 32]       29,750\n",
       "│    │    └─MBConvBlock: 3-6                            [1, 40, 32, 32, 32]       55,290\n",
       "│    │    └─MBConvBlock: 3-7                            [1, 80, 16, 16, 16]       41,450\n",
       "│    │    └─MBConvBlock: 3-8                            [1, 80, 16, 16, 16]       111,540\n",
       "│    │    └─MBConvBlock: 3-9                            [1, 80, 16, 16, 16]       111,540\n",
       "│    │    └─MBConvBlock: 3-10                           [1, 112, 16, 16, 16]      174,004\n",
       "│    │    └─MBConvBlock: 3-11                           [1, 112, 16, 16, 16]      275,772\n",
       "│    │    └─MBConvBlock: 3-12                           [1, 112, 16, 16, 16]      275,772\n",
       "│    │    └─MBConvBlock: 3-13                           [1, 192, 8, 8, 8]         329,692\n",
       "│    │    └─MBConvBlock: 3-14                           [1, 192, 8, 8, 8]         703,152\n",
       "│    │    └─MBConvBlock: 3-15                           [1, 192, 8, 8, 8]         703,152\n",
       "│    │    └─MBConvBlock: 3-16                           [1, 192, 8, 8, 8]         703,152\n",
       "│    │    └─MBConvBlock: 3-17                           [1, 320, 8, 8, 8]         737,968\n",
       "├─UnetDecoder: 1-2                                      [1, 16, 256, 256, 256]    --\n",
       "│    └─Identity: 2-5                                    [1, 320, 8, 8, 8]         --\n",
       "│    └─ModuleList: 2-6                                  --                        --\n",
       "│    │    └─DecoderBlock: 3-18                          [1, 256, 16, 16, 16]      4,756,480\n",
       "│    │    └─DecoderBlock: 3-19                          [1, 128, 32, 32, 32]      1,465,856\n",
       "│    │    └─DecoderBlock: 3-20                          [1, 64, 64, 64, 64]       373,504\n",
       "│    │    └─DecoderBlock: 3-21                          [1, 32, 128, 128, 128]    110,720\n",
       "│    │    └─DecoderBlock: 3-22                          [1, 16, 256, 256, 256]    20,800\n",
       "├─SegmentationHead: 1-3                                 [1, 3, 256, 256, 256]     --\n",
       "│    └─Conv3d: 2-7                                      [1, 3, 256, 256, 256]     1,299\n",
       "│    └─Identity: 2-8                                    [1, 3, 256, 256, 256]     --\n",
       "│    └─Activation: 2-9                                  [1, 3, 256, 256, 256]     --\n",
       "│    │    └─Identity: 3-23                              [1, 3, 256, 256, 256]     --\n",
       "=========================================================================================================\n",
       "Total params: 11,418,767\n",
       "Trainable params: 11,418,767\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.GIGABYTES): 766.95\n",
       "=========================================================================================================\n",
       "Input size (MB): 201.33\n",
       "Forward/backward pass size (MB): 16529.10\n",
       "Params size (MB): 27.07\n",
       "Estimated Total Size (MB): 16757.50\n",
       "========================================================================================================="
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import segmentation_models_pytorch_3d as smp\n",
    "from segmentation_models_pytorch_3d.encoders import get_preprocessing_fn\n",
    "from torchinfo import summary\n",
    "\n",
    "net = smp.Unet(\n",
    "    encoder_name=\"efficientnet-b0\", # choose encoder, e.g. resnet34\n",
    "    in_channels=3,                  # model input channels (1 for gray-scale volumes, 3 for RGB, etc.)\n",
    "    classes=3,                      # model output channels (number of classes in your dataset)\n",
    ")\n",
    "\n",
    "# Preparing your data the same way as during weights pre-training may give you better results\n",
    "preprocess_input = get_preprocessing_fn('resnet34', pretrained='imagenet')\n",
    "print(preprocess_input)\n",
    "summary(net,(1,3,256,256,256)) # batch size, channels, height, width, depth\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## models from Segmentation Models Pytroch 2D\n",
    "* Reference: https://smp.readthedocs.io/en/latest/index.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "functools.partial(<function preprocess_input at 0x0000016E9B684040>, input_space='RGB', input_range=[0, 1], mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "===============================================================================================\n",
       "Layer (type:depth-idx)                        Output Shape              Param #\n",
       "===============================================================================================\n",
       "Unet                                          [1, 3, 256, 256]          --\n",
       "├─ResNetEncoder: 1-1                          [1, 3, 256, 256]          --\n",
       "│    └─Conv2d: 2-1                            [1, 64, 128, 128]         9,408\n",
       "│    └─BatchNorm2d: 2-2                       [1, 64, 128, 128]         128\n",
       "│    └─ReLU: 2-3                              [1, 64, 128, 128]         --\n",
       "│    └─MaxPool2d: 2-4                         [1, 64, 64, 64]           --\n",
       "│    └─Sequential: 2-5                        [1, 64, 64, 64]           --\n",
       "│    │    └─BasicBlock: 3-1                   [1, 64, 64, 64]           73,984\n",
       "│    │    └─BasicBlock: 3-2                   [1, 64, 64, 64]           73,984\n",
       "│    │    └─BasicBlock: 3-3                   [1, 64, 64, 64]           73,984\n",
       "│    └─Sequential: 2-6                        [1, 128, 32, 32]          --\n",
       "│    │    └─BasicBlock: 3-4                   [1, 128, 32, 32]          230,144\n",
       "│    │    └─BasicBlock: 3-5                   [1, 128, 32, 32]          295,424\n",
       "│    │    └─BasicBlock: 3-6                   [1, 128, 32, 32]          295,424\n",
       "│    │    └─BasicBlock: 3-7                   [1, 128, 32, 32]          295,424\n",
       "│    └─Sequential: 2-7                        [1, 256, 16, 16]          --\n",
       "│    │    └─BasicBlock: 3-8                   [1, 256, 16, 16]          919,040\n",
       "│    │    └─BasicBlock: 3-9                   [1, 256, 16, 16]          1,180,672\n",
       "│    │    └─BasicBlock: 3-10                  [1, 256, 16, 16]          1,180,672\n",
       "│    │    └─BasicBlock: 3-11                  [1, 256, 16, 16]          1,180,672\n",
       "│    │    └─BasicBlock: 3-12                  [1, 256, 16, 16]          1,180,672\n",
       "│    │    └─BasicBlock: 3-13                  [1, 256, 16, 16]          1,180,672\n",
       "│    └─Sequential: 2-8                        [1, 512, 8, 8]            --\n",
       "│    │    └─BasicBlock: 3-14                  [1, 512, 8, 8]            3,673,088\n",
       "│    │    └─BasicBlock: 3-15                  [1, 512, 8, 8]            4,720,640\n",
       "│    │    └─BasicBlock: 3-16                  [1, 512, 8, 8]            4,720,640\n",
       "├─UnetDecoder: 1-2                            [1, 16, 256, 256]         --\n",
       "│    └─Identity: 2-9                          [1, 512, 8, 8]            --\n",
       "│    └─ModuleList: 2-10                       --                        --\n",
       "│    │    └─DecoderBlock: 3-17                [1, 256, 16, 16]          2,360,320\n",
       "│    │    └─DecoderBlock: 3-18                [1, 128, 32, 32]          590,336\n",
       "│    │    └─DecoderBlock: 3-19                [1, 64, 64, 64]           147,712\n",
       "│    │    └─DecoderBlock: 3-20                [1, 32, 128, 128]         46,208\n",
       "│    │    └─DecoderBlock: 3-21                [1, 16, 256, 256]         6,976\n",
       "├─SegmentationHead: 1-3                       [1, 3, 256, 256]          --\n",
       "│    └─Conv2d: 2-11                           [1, 3, 256, 256]          435\n",
       "│    └─Identity: 2-12                         [1, 3, 256, 256]          --\n",
       "│    └─Activation: 2-13                       [1, 3, 256, 256]          --\n",
       "│    │    └─Identity: 3-22                    [1, 3, 256, 256]          --\n",
       "===============================================================================================\n",
       "Total params: 24,436,659\n",
       "Trainable params: 24,436,659\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.GIGABYTES): 7.83\n",
       "===============================================================================================\n",
       "Input size (MB): 0.79\n",
       "Forward/backward pass size (MB): 144.70\n",
       "Params size (MB): 97.75\n",
       "Estimated Total Size (MB): 243.24\n",
       "==============================================================================================="
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import segmentation_models_pytorch as smp\n",
    "from segmentation_models_pytorch.encoders import get_preprocessing_fn\n",
    "\n",
    "from torchinfo import summary\n",
    "\n",
    "# Unet Unet++ FPN PSPNet DeepLabV3 DeepLabV3+ Linknet MAnet PAN\n",
    "\n",
    "net = smp.Unet(\n",
    "    encoder_name=\"resnet34\",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7\n",
    "    encoder_weights=\"imagenet\",     # use `imagenet` pre-trained weights for encoder initialization\n",
    "    in_channels=3,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)\n",
    "    classes=3,                      # model output channels (number of classes in your dataset)\n",
    ")\n",
    "\n",
    "# Preparing your data the same way as during weights pre-training may give you better results\n",
    "preprocess_input = get_preprocessing_fn('resnet34', pretrained='imagenet')\n",
    "print(preprocess_input)\n",
    "summary(net,(1,3,256,256))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.11.6 ('base')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "d0af3871f74fa798a0a222812f3db7ca3ddb43c51379edde28baf2323a73ac20"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}