Powered by AppSignal & Oban Pro
Would you like to see your link here? Contact us
Notesclub

Don't Repeat Yourself

dont-repeat-yourself.livemd

Don’t Repeat Yourself

Mix.install([
  {:axon_onnx, git: "https://github.com/mortont/axon_onnx.git"},
  {:axon, "~> 0.5"},
  {:nx, "~> 0.5"},
  {:exla, "~> 0.5"},
  {:stb_image, "~> 0.6"},
  {:kino, "~> 0.8"}
])

Cats and Dogs Again

Application.put_env(
  :exla,
  :clients,
  cuda: [
    platforms: :cuda,
    lazy_transformers: :never
  ]
)

Nx.global_default_backend(EXLA.Backend)
Nx.Defn.default_options(compiler: EXLA)
[]
defmodule CatsAndDogs do
  def pipeline(paths, batch_size, target_height, target_width) do
    paths
    |> Enum.shuffle()
    |> Task.async_stream(&parse_image/1)
    |> Stream.filter(fn
      {:ok, {%StbImage{}, _}} -> true
      _ -> false
    end)
    |> Stream.map(&to_tensors(&1, target_height, target_width))
    |> Stream.chunk_every(batch_size, batch_size, :discard)
    |> Stream.map(fn chunks ->
      {img_chunk, label_chunk} = Enum.unzip(chunks)
      {Nx.stack(img_chunk), Nx.stack(label_chunk)}
    end)
  end

  def pipeline_with_augmentations(
        paths,
        batch_size,
        target_height,
        target_width
      ) do
    paths
    |> Enum.shuffle()
    |> Task.async_stream(&parse_image/1)
    |> Stream.filter(fn
      {:ok, {%StbImage{}, _}} -> true
      _ -> false
    end)
    |> Stream.map(&to_tensors(&1, target_height, target_width))
    |> Stream.map(&random_flip(&1, :height))
    |> Stream.map(&random_flip(&1, :width))
    |> Stream.chunk_every(batch_size, batch_size, :discard)
    |> Stream.map(fn chunks ->
      {img_chunk, label_chunk} = Enum.unzip(chunks)
      {Nx.stack(img_chunk), Nx.stack(label_chunk)}
    end)
  end

  defp parse_image(path) do
    label = if String.contains?(path, "cat"), do: 0, else: 1

    case StbImage.read_file(path) do
      {:ok, img} -> {img, label}
      _error -> :error
    end
  end

  defp to_tensors({:ok, {img, label}}, target_height, target_width) do
    img_tensor =
      img
      |> StbImage.resize(target_height, target_width)
      |> StbImage.to_nx()
      |> Nx.divide(255)
      |> Nx.transpose(axes: [:channels, :height, :width])

    label_tensor = Nx.tensor([label])
    {img_tensor, label_tensor}
  end

  defp random_flip({image, label}, axis) do
    if :rand.uniform() < 0.5 do
      {Nx.reverse(image, axes: [axis]), label}
    else
      {image, label}
    end
  end
end
{:module, CatsAndDogs, <<70, 79, 82, 49, 0, 0, 20, ...>>, {:random_flip, 2}}
{test_paths, train_paths} =
  Path.wildcard("train/*.jpg")
  |> Enum.shuffle()
  |> Enum.split(1000)

{test_paths, val_paths} = Enum.split(test_paths, 750)

batch_size = 32
target_height = 160
target_width = 160

train_pipeline =
  CatsAndDogs.pipeline_with_augmentations(
    train_paths,
    batch_size,
    target_height,
    target_width
  )

val_pipeline =
  CatsAndDogs.pipeline(
    val_paths,
    batch_size,
    target_height,
    target_width
  )

test_pipeline =
  CatsAndDogs.pipeline(
    test_paths,
    batch_size,
    target_height,
    target_width
  )

Enum.take(train_pipeline, 1)

18:00:37.908 [debug] No platform configuration specified, falling back to host platform
Available platforms are: [:host, :interpreter]

[
  {#Nx.Tensor<
     f32[32][channels: 3][height: 160][width: 160]
     EXLA.Backend
     [
       [
         [
           [0.2980392277240753, 0.29411765933036804, 0.27843138575553894, 0.27450981736183167, 0.27843138575553894, 0.2823529541492462, 0.2862745225429535, 0.2862745225429535, 0.2862745225429535, 0.2862745225429535, 0.2862745225429535, 0.29019609093666077, 0.29019609093666077, 0.29019609093666077, 0.29019609093666077, 0.29019609093666077, 0.2862745225429535, 0.2862745225429535, 0.2862745225429535, 0.2823529541492462, 0.2862745225429535, 0.29411765933036804, 0.3019607961177826, 0.2980392277240753, 0.29411765933036804, 0.29411765933036804, 0.29411765933036804, 0.2980392277240753, 0.30588236451148987, 0.30588236451148987, 0.30588236451148987, 0.30588236451148987, 0.30588236451148987, 0.30980393290519714, 0.3176470696926117, 0.3137255012989044, 0.3137255012989044, 0.32156863808631897, 0.30980393290519714, 0.3137255012989044, 0.32156863808631897, 0.3176470696926117, 0.32156863808631897, 0.32549020648002625, 0.3294117748737335, 0.34117648005485535, 0.3490196168422699, 0.3490196168422699, ...],
           ...
         ],
         ...
       ],
       ...
     ]
   >,
   #Nx.Tensor<
     s64[32][1]
     EXLA.Backend
     [
       [0],
       [0],
       [1],
       [0],
       [1],
       [0],
       [1],
       [0],
       [1],
       [1],
       [1],
       [1],
       [1],
       [1],
       [0],
       [0],
       [1],
       [1],
       [1],
       [0],
       [1],
       [1],
       [0],
       [0],
       [1],
       [0],
       [1],
       [0],
       [0],
       [0],
       [0],
       [1]
     ]
   >}
]
{cnn_base, cnn_base_params} = AxonOnnx.import("mobilenetv2-7.onnx", batch_size: batch_size)
...
     >
   },
   "mobilenetv20_features_linearbottleneck9_conv2_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[64][384][1][1]
       EXLA.Backend
       [
         [
           [
             [-0.029621589928865433]
           ],
           [
             [0.03196175768971443]
           ],
           [
             [-0.04775437340140343]
           ],
           [
             [-0.02047593705356121]
           ],
           [
             [0.022681893780827522]
           ],
           [
             [0.0061656758189201355]
           ],
           [
             [-0.020124075934290886]
           ],
           [
             [-0.0419141948223114]
           ],
           [
             [-0.0958712100982666]
           ],
           [
             [-0.03112773969769478]
           ],
           [
             [0.033005211502313614]
           ],
           [
             [-0.0796334370970726]
           ],
           [
             [-0.08826649934053421]
           ],
           [
             [0.008132039569318295]
           ],
           [
             [0.13005898892879486]
           ],
           [
             [0.01112390961498022]
           ],
           [
             [-0.025056127458810806]
           ],
           [
             [-0.021656425669789314]
           ],
           [
             [-0.009619093500077724]
           ],
           [
             [-0.09488093852996826]
           ],
           [
             [-0.03190577030181885]
           ],
           [
             [0.08531086146831512]
           ],
           [
             [0.047572940587997437]
           ],
           [
             [-0.036207523196935654]
           ],
           [
             [-0.1315651535987854]
           ],
           [
             [-0.03746785223484039]
           ],
           [
             [-0.04043968394398689]
           ],
           [
             [-0.1008371040225029]
           ],
           [
             [-7.86909949965775e-4]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck11_conv2_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[96][576][1][1]
       EXLA.Backend
       [
         [
           [
             [-0.04216459393501282]
           ],
           [
             [-0.03359765186905861]
           ],
           [
             [0.021302534267306328]
           ],
           [
             [0.06966963410377502]
           ],
           [
             [-0.023086098954081535]
           ],
           [
             [0.02423691935837269]
           ],
           [
             [-0.0027191585395485163]
           ],
           [
             [-0.12812955677509308]
           ],
           [
             [0.10142464935779572]
           ],
           [
             [0.04367773234844208]
           ],
           [
             [-0.04808792844414711]
           ],
           [
             [-0.059124793857336044]
           ],
           [
             [-0.03544275462627411]
           ],
           [
             [-0.10292188823223114]
           ],
           [
             [-0.15540222823619843]
           ],
           [
             [0.08993023633956909]
           ],
           [
             [-0.0018712839810177684]
           ],
           [
             [0.04499407112598419]
           ],
           [
             [0.03717062994837761]
           ],
           [
             [0.1088862419128418]
           ],
           [
             [0.03789613023400307]
           ],
           [
             [0.011467562057077885]
           ],
           [
             [0.022153226658701897]
           ],
           [
             [0.011087901890277863]
           ],
           [
             [0.10830669850111008]
           ],
           [
             [-0.039926622062921524]
           ],
           [
             [0.03126271069049835]
           ],
           [
             [0.11330234259366989]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck2_batchnorm1_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[144]
       EXLA.Backend
       [0.06266669183969498, 0.34174367785453796, 0.014498202130198479, 0.09171395748853683, -0.01248596515506506, 0.09248802065849304, 0.03913332521915436, -0.12479757517576218, 0.04716098681092262, 0.2623796761035919, 0.00875355675816536, 0.3866884410381317, 0.34793952107429504, 0.07759835571050644, 0.28590700030326843, 0.34494516253471375, 0.1602264642715454, 0.3233429491519928, -0.08500878512859344, 0.017791876569390297, 0.012547487393021584, 0.07443935424089432, 0.470037043094635, 0.06998748332262039, -0.1145705133676529, 0.35432684421539307, 0.0036894683726131916, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[144]
       EXLA.Backend
       [0.14569389820098877, 0.08483350276947021, 0.18709848821163177, 0.08500487357378006, 0.0695008710026741, 0.13266052305698395, 0.23066304624080658, 0.1294335573911667, 0.15951289236545563, 0.1612233966588974, 0.04986649379134178, 0.24220621585845947, 0.13059896230697632, 0.2140040248632431, 0.16401590406894684, 0.04375945031642914, 0.1284344643354416, 0.10862768441438675, 0.22237107157707214, 0.20400500297546387, 0.2341395914554596, 0.36944401264190674, 0.1626269817352295, 0.1425088793039322, 0.37626516819000244, 0.10176140815019608, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[144]
       EXLA.Backend
       [0.04842672124505043, 0.011466434225440025, -0.28272300958633423, 0.04086904227733612, -5.605193857299268e-45, 0.03728378936648369, -0.07810080051422119, -0.0095414062961936, 0.019914837554097176, -0.26291584968566895, 1.200515914815965e-22, 0.4378272294998169, -0.07176417112350464, -0.20427647233009338, -0.40337827801704407, 0.07497201859951019, 0.06511131674051285, 0.02972135879099369, -0.029825769364833832, 0.03816595301032066, 0.008758057840168476, -0.5172790288925171, 0.18057824671268463, 0.03219117224216461, 0.25589102506637573, ...]
     >,
     "var" => #Nx.Tensor<
       f32[144]
       EXLA.Backend
       [0.0016082772053778172, 0.008874273858964443, 0.014802338555455208, 0.010124930180609226, 5.605193857299268e-45, 0.0027794779743999243, 0.019753890112042427, 6.709994049742818e-4, 0.004814377520233393, 0.018505984917283058, 3.875728477835647e-26, 0.024166392162442207, 0.011687243357300758, 0.018400687724351883, 0.00805666297674179, 0.003925488330423832, 0.004267431795597076, 0.005561113357543945, 0.00586837250739336, 0.01074590440839529, 0.023998573422431946, 0.010233975946903229, 0.007847330532968044, 0.0024662334471940994, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck3_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[144][1][3][3]
       EXLA.Backend
       [
         [
           [
             [0.007945280522108078, -0.5135775208473206, 0.642527163028717],
             [0.09571424126625061, -0.4191504120826721, -0.42046400904655457],
             [0.2738749086856842, -1.341597080230713, 1.7941224575042725]
           ]
         ],
         [
           [
             [0.36698344349861145, 0.8006415963172913, 1.0355921983718872],
             [0.2570298910140991, 0.9058852791786194, 0.1435706615447998],
             [0.08509054780006409, -0.29565608501434326, -0.19910597801208496]
           ]
         ],
         [
           [
             [-0.07826036959886551, 1.5095280408859253, -0.7258926033973694],
             [-0.1067059189081192, -0.07988816499710083, -0.4704415500164032],
             [-2.0861942768096924, 0.03404847905039787, ...]
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck10_batchnorm0_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[384]
       EXLA.Backend
       [-0.14200253784656525, -0.08138399571180344, 0.020918697118759155, -0.030422303825616837, -0.054856158792972565, 0.0021326385904103518, -0.18490885198116302, -0.03978404402732849, 0.17561353743076324, 0.007802810054272413, -0.05954292416572571, 0.001685986528173089, -0.021205948665738106, 0.09003060311079025, 0.03726910054683685, 0.14319199323654175, 0.15906931459903717, 0.09075694531202316, 0.14122611284255981, 0.19317775964736938, 0.11409415304660797, 0.2862238883972168, 0.21863330900669098, 0.31881415843963623, 0.15870404243469238, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[384]
       EXLA.Backend
       [0.08955207467079163, 0.12251311540603638, 0.13474930822849274, 0.12169622629880905, 0.10408146679401398, 0.12673817574977875, 0.09429759532213211, 0.004978157579898834, 0.06450905650854111, 0.12292398512363434, 0.008911419659852982, 0.16245806217193604, 0.12929747998714447, 0.07689117640256882, 0.13534259796142578, 0.0939403846859932, 0.07975732535123825, 0.07448924332857132, 0.02857014909386635, 0.05634821206331253, 0.04961058497428894, 0.08511262387037277, 0.1229783222079277, 0.09482692182064056, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[384]
       EXLA.Backend
       [2.3076812794897705e-4, 8.846540004014969e-4, 2.087028551613912e-4, 2.626747591421008e-4, -1.6135886835400015e-4, 3.3246324164792895e-4, 1.477891200920567e-4, 8.788402919890359e-5, -5.301489727571607e-4, 3.370658669155091e-4, 1.6383620095439255e-4, 8.14800092484802e-4, 4.6927749644964933e-4, 2.0585938182193786e-4, 7.599047385156155e-4, -1.2353641795925796e-4, -2.2819229343440384e-4, 1.4898847439326346e-4, -1.760231316438876e-5, 1.1794252350227907e-4, 3.7943560164421797e-4, 8.82304084370844e-6, 8.040065877139568e-4, ...]
     >,
     "var" => #Nx.Tensor<
       f32[384]
       EXLA.Backend
       [0.03708014264702797, 0.1853957325220108, 0.18863092362880707, 0.12991510331630707, 0.08043266087770462, 0.15097717940807343, 0.07633852958679199, 0.003970237914472818, 0.38535448908805847, 0.13255122303962708, 0.01239611953496933, 0.37033578753471375, 0.1896938532590866, 0.13335467875003815, 0.267678827047348, 0.12042298913002014, 0.16346299648284912, 0.09761403501033783, 0.12983210384845734, 0.211727574467659, 0.12430556863546371, 0.3960045278072357, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck14_batchnorm2_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[160]
       EXLA.Backend
       [-8.388164474126825e-8, -5.7530677111117257e-8, 7.027283999150313e-8, -9.3144599588868e-8, 4.638787132194011e-8, -9.98987204070545e-8, 1.4243799739688257e-7, -1.1660619492204205e-7, 6.76579361424956e-8, 1.2566950147174794e-7, -1.1262851984383815e-7, 6.708535948973804e-8, 2.0174957171548158e-7, 1.1858566040245933e-7, 3.9078233982081656e-8, -1.565037024420235e-7, 1.173001891174863e-7, -1.3883133931358316e-7, -1.3765874484761298e-7, -1.1415199452358138e-7, -3.374253765286994e-8, 1.203010810968408e-7, -1.509711751168652e-7, 6.394230922524002e-8, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[160]
       EXLA.Backend
       [0.18752147257328033, 0.34522393345832825, 0.24475780129432678, 0.33060014247894287, 0.21970151364803314, 0.3153970539569855, 0.28512927889823914, 0.169156014919281, 0.3166297674179077, 0.2279355376958847, 0.21180318295955658, 0.22759808599948883, 0.23080326616764069, 0.17287319898605347, 0.27370786666870117, 0.2255641222000122, 0.2651505172252655, 0.17854700982570648, 0.257531076669693, 0.2266509234905243, 0.2487022578716278, 0.22212600708007812, 0.19973145425319672, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[160]
       EXLA.Backend
       [-0.10289469361305237, 0.13421736657619476, 0.3127436637878418, -0.010098671540617943, 0.11501706391572952, 0.07170626521110535, -0.12338873744010925, -0.09662780910730362, 0.10249986499547958, 0.039654918015003204, 0.23144079744815826, -0.02762533910572529, -0.2323276251554489, -0.14802084863185883, -0.17405559122562408, -0.05928152799606323, 0.30801624059677124, 0.03448152542114258, 0.10294467210769653, -0.14233750104904175, 0.042918454855680466, -0.15366889536380768, ...]
     >,
     "var" => #Nx.Tensor<
       f32[160]
       EXLA.Backend
       [0.01982937380671501, 0.07223469018936157, 0.03559407591819763, 0.05254514887928963, 0.03399147093296051, 0.05125325173139572, 0.045940618962049484, 0.015112430788576603, 0.0665295273065567, 0.029054192826151848, 0.024110015481710434, 0.028089039027690887, 0.032291218638420105, 0.017801793292164803, 0.04934801161289215, 0.028767650946974754, 0.036390747874975204, 0.020782455801963806, 0.038260187953710556, 0.03192899003624916, 0.0354829803109169, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck0_batchnorm2_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[16]
       EXLA.Backend
       [-0.0010560675291344523, -6.787155289202929e-4, -7.046025712043047e-4, -9.398604743182659e-4, -8.966192253865302e-4, -2.11737904464826e-4, -5.149749922566116e-4, -0.0010324727045372128, -7.925425306893885e-4, -9.090156527236104e-5, -4.453391011338681e-4, -9.805577574297786e-4, -7.45584664400667e-4, -7.640189724043012e-4, -0.0011470100143924356, -3.8889769348315895e-4]
     >,
     "gamma" => #Nx.Tensor<
       f32[16]
       EXLA.Backend
       [0.5097491145133972, 0.6192775368690491, 0.657399594783783, 0.5558041930198669, 0.6164559721946716, 0.6381960511207581, 0.6152222156524658, 0.4946329891681671, 0.5616724491119385, 0.6567564606666565, 0.5526584982872009, 0.42007970809936523, 0.4637764096260071, 0.5626768469810486, 0.6189385056495667, 0.6354565620422363]
     >,
     "mean" => #Nx.Tensor<
       f32[16]
       EXLA.Backend
       [0.020797766745090485, -0.4769341051578522, -0.4614068269729614, 1.7337007522583008, 0.3933316469192505, -0.8451346755027771, 1.245340347290039, -0.48207274079322815, 0.8380969166755676, -2.1938257217407227, 2.286595106124878, 0.04284362867474556, -1.1523691415786743, -0.37927713990211487, -3.048325300216675, -0.7846338748931885]
     >,
     "var" => #Nx.Tensor<
       f32[16]
       EXLA.Backend
       [0.020222526043653488, 0.055290378630161285, 0.06333210319280624, 0.026319650933146477, 0.03692027926445007, 0.0943082720041275, 0.044809967279434204, 0.05162627622485161, 0.03870456665754318, 0.13497591018676758, 0.04896015673875809, 0.0536123663187027, 0.03798253461718559, 0.045080412179231644, 0.05210200324654579, 0.049917370080947876]
     >
   },
   "mobilenetv20_features_linearbottleneck7_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[384][1][3][3]
       EXLA.Backend
       [
         [
           [
             [0.09069628268480301, 0.011556080542504787, -0.044405367225408554],
             [0.1974998563528061, -0.05144780874252319, -0.15406860411167145],
             [0.03294772282242775, -0.08055609464645386, -0.10401368886232376]
           ]
         ],
         [
           [
             [0.004530614707618952, -0.016135336831212044, 0.010837683454155922],
             [-0.013911036774516106, -0.17941752076148987, 0.0742872878909111],
             [0.03701423481106758, 0.17489252984523773, 0.049539677798748016]
           ]
         ],
         [
           [
             [-0.03442509472370148, -0.04338786378502846, -8.070958429016173e-4],
             [-0.06931383162736893, ...],
             ...
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck1_conv0_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[96][16][1][1]
       EXLA.Backend
       [
         [
           [
             [-0.26627230644226074]
           ],
           [
             [0.1223592683672905]
           ],
           [
             [-0.17675793170928955]
           ],
           [
             [-0.14175795018672943]
           ],
           [
             [-0.04847275838255882]
           ],
           [
             [0.02572016417980194]
           ],
           [
             [0.18247880041599274]
           ],
           [
             [-0.22272470593452454]
           ],
           [
             [-0.0375950001180172]
           ],
           [
             [-0.12036115676164627]
           ],
           [
             [-0.42806223034858704]
           ],
           [
             [0.2049194574356079]
           ],
           [
             [0.025416148826479912]
           ],
           [
             [-0.20868347585201263]
           ],
           [
             [-0.00869692862033844]
           ],
           [
             [0.10014521330595016]
           ]
         ],
         [
           [
             [0.1901395469903946]
           ],
           [
             [-0.15854698419570923]
           ],
           [
             [-0.05858785659074783]
           ],
           [
             [-0.016011832281947136]
           ],
           [
             [-0.07135247439146042]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck2_batchnorm2_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[24]
       EXLA.Backend
       [-6.408389890566468e-5, -1.0128262510988861e-4, -4.3479518353706226e-5, -1.589354214956984e-4, -7.557999197160825e-5, -2.3196584152174182e-5, -2.216532493548584e-6, 9.525781933916733e-6, -8.600373985245824e-5, -2.3008385323919356e-4, -2.3174902889877558e-4, -2.012878976529464e-4, -9.838923870120198e-5, -1.22729703434743e-4, -9.646145917940885e-5, -2.797471643134486e-5, -1.039204653352499e-4, -1.2249314750079066e-4, -8.16335596027784e-5, -1.541685633128509e-4, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[24]
       EXLA.Backend
       [0.5065965056419373, 0.44786974787712097, 0.5534567832946777, 0.3781253695487976, 0.44674113392829895, 0.28593751788139343, 0.5029821395874023, 0.19650447368621826, 0.3918337821960449, 0.33050450682640076, 0.31709131598472595, 0.1402808576822281, 0.4198921024799347, 0.02611447498202324, 0.5337468981742859, 0.49816271662712097, 0.32386210560798645, 0.40961867570877075, 0.3363959491252899, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[24]
       EXLA.Backend
       [-0.49676281213760376, -0.5412184000015259, 0.2777172029018402, 0.30402812361717224, -0.038367077708244324, 0.19269412755966187, -0.20861922204494476, 0.3330950438976288, -0.35258007049560547, -0.0413212850689888, -0.44278526306152344, 0.43556883931159973, -0.007557356264442205, -0.049789782613515854, -0.1810888946056366, -0.1666330248117447, 0.15477608144283295, 0.1351827085018158, ...]
     >,
     "var" => #Nx.Tensor<
       f32[24]
       EXLA.Backend
       [0.0421639047563076, 0.03801903501152992, 0.04582618921995163, 0.028266852721571922, 0.025787167251110077, 0.026369646191596985, 0.03499341756105423, 0.02089923992753029, 0.02381058968603611, 0.04512317106127739, 0.018083490431308746, 0.02251904457807541, 0.023619014769792557, 0.005619029048830271, 0.04942920058965683, 0.02654757723212242, 0.02269740216434002, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck8_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[384][1][3][3]
       EXLA.Backend
       [
         [
           [
             [-0.0020327651873230934, -0.03829871118068695, -0.013563335873186588],
             [-0.07023981958627701, -0.0530773401260376, -0.05999664217233658],
             [-0.02172546088695526, -0.05160605534911156, -0.02340804785490036]
           ]
         ],
         [
           [
             [-0.06414368003606796, -0.09891236573457718, 0.09147416055202484],
             [-0.035863131284713745, -0.026917265728116035, 0.07736452668905258],
             [0.08244364708662033, 0.10789858549833298, -0.16784991323947906]
           ]
         ],
         [
           [
             [0.13550637662410736, ...],
             ...
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck11_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[576][1][3][3]
       EXLA.Backend
       [
         [
           [
             [0.0561985969543457, -0.08534234762191772, -0.0728735700249672],
             [0.20427849888801575, -0.03384411707520485, -0.18976956605911255],
             [0.10962294042110443, 0.006988390814512968, -0.03254517540335655]
           ]
         ],
         [
           [
             [-0.007281317375600338, -0.10424316674470901, -0.0011419359361752868],
             [-0.0659799799323082, -0.2147035449743271, -0.10787361860275269],
             [0.004407626111060381, 0.44921964406967163, 0.06901120394468307]
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck14_conv0_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[960][160][1][1]
       EXLA.Backend
       [
         [
           [
             [0.037727002054452896]
           ],
           [
             [-0.07920914888381958]
           ],
           [
             [-0.0031863711774349213]
           ],
           [
             [-0.017933860421180725]
           ],
           [
             [-0.012083934620022774]
           ],
           [
             [-0.03305855393409729]
           ],
           [
             [0.010104630142450333]
           ],
           [
             [-0.01614421419799328]
           ],
           [
             [0.0415346585214138]
           ],
           [
             [-0.05910014733672142]
           ],
           [
             [3.380556881893426e-4]
           ],
           [
             [0.02735193446278572]
           ],
           [
             [-0.0788668617606163]
           ],
           [
             [-0.028101010248064995]
           ],
           [
             [-0.01531447283923626]
           ],
           [
             [-0.04052167013287544]
           ],
           [
             [0.0598626434803009]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck15_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[960][1][3][3]
       EXLA.Backend
       [
         [
           [
             [0.008491688407957554, 0.01882096938788891, -0.001548346015624702],
             [0.009763970039784908, 0.2359771430492401, 0.03868749365210533],
             [-0.05791253224015236, 0.028106149286031723, -0.07315897196531296]
           ]
         ],
         [
           [
             [-0.01212818082422018, -0.16292721033096313, -0.033156707882881165],
             [0.005929793696850538, 0.2561517059803009, 0.014003193937242031],
             [-0.04701530933380127, ...]
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck7_conv2_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[64][384][1][1]
       EXLA.Backend
       [
         [
           [
             [-0.03408721089363098]
           ],
           [
             [-0.017362290993332863]
           ],
           [
             [0.01995295286178589]
           ],
           [
             [-0.051705457270145416]
           ],
           [
             [0.0689198449254036]
           ],
           [
             [0.08685684949159622]
           ],
           [
             [-0.09933704882860184]
           ],
           [
             [0.07178854197263718]
           ],
           [
             [0.044526729732751846]
           ],
           [
             [-0.018590005114674568]
           ],
           [
             [0.013196618296205997]
           ],
           [
             [-0.0795169249176979]
           ],
           [
             [-0.10800802707672119]
           ],
           [
             [-0.042785294353961945]
           ],
           [
             [0.028526680544018745]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck4_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[192][1][3][3]
       EXLA.Backend
       [
         [
           [
             [-0.031621839851140976, 0.18290314078330994, 0.040609210729599],
             [-0.19228996336460114, 0.031516727060079575, 0.09460115432739258],
             [-0.026457147672772408, 0.05563095211982727, 0.0026122310664504766]
           ]
         ],
         [
           [
             [-0.00655049504712224, -0.004564541857689619, -0.01380461547523737],
             [-0.002132135210558772, -0.17951297760009766, ...],
             ...
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck16_conv0_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[960][160][1][1]
       EXLA.Backend
       [
         [
           [
             [0.04804954305291176]
           ],
           [
             [-0.0361880324780941]
           ],
           [
             [0.09869731217622757]
           ],
           [
             [0.03218498453497887]
           ],
           [
             [-0.05613848939538002]
           ],
           [
             [0.03235843405127525]
           ],
           [
             [0.035072293132543564]
           ],
           [
             [0.04535597562789917]
           ],
           [
             [0.031406670808792114]
           ],
           [
             [0.003971233498305082]
           ],
           [
             [0.016072509810328484]
           ],
           [
             [-0.041956160217523575]
           ],
           [
             [-0.018940964713692665]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck10_batchnorm2_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[96]
       EXLA.Backend
       [2.1335902999908285e-7, 3.535939754328865e-7, 5.073925990473072e-7, 5.08670098042785e-7, 3.3390904263796983e-7, 4.946827516505437e-7, 4.795706445293035e-7, 3.743390948329761e-7, 5.272184466775798e-7, 5.93354513966915e-7, 2.1419212714590685e-7, 6.500247309304541e-7, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[96]
       EXLA.Backend
       [0.2979993522167206, 0.3140060007572174, 0.2445983588695526, 0.2865447998046875, 0.35691216588020325, 0.3592449426651001, 0.27936986088752747, 0.25922414660453796, 0.2633695900440216, 0.2901868522167206, 0.2989289462566376, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[96]
       EXLA.Backend
       [-0.02239207923412323, -0.5143733024597168, 0.05194510519504547, -0.3102657198905945, -0.03491494804620743, 0.21602749824523926, 0.09913935512304306, -0.2323945015668869, -0.47182056307792664, 0.019973447546362877, ...]
     >,
     "var" => #Nx.Tensor<
       f32[96]
       EXLA.Backend
       [0.10519497841596603, 0.09767355769872665, 0.06126777082681656, 0.09396423399448395, 0.17368470132350922, 0.14595730602741241, 0.0813060998916626, 0.07312573492527008, 0.059541765600442886, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck14_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[960][1][3][3]
       EXLA.Backend
       [
         [
           [
             [0.03072543255984783, 0.024941517040133476, -0.044987138360738754],
             [0.2105334848165512, -0.09797866642475128, -0.12514935433864594],
             [0.049303051084280014, -0.004124515224248171, -0.0421803742647171]
           ]
         ],
         [
           [
             [0.027953263372182846, -0.0323273129761219, ...],
             ...
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck12_conv0_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[576][96][1][1]
       EXLA.Backend
       [
         [
           [
             [-0.020311787724494934]
           ],
           [
             [0.00342985475435853]
           ],
           [
             [-0.04346230626106262]
           ],
           [
             [-0.016626032069325447]
           ],
           [
             [-0.012114123441278934]
           ],
           [
             [-0.012785198166966438]
           ],
           [
             [0.056096695363521576]
           ],
           [
             [0.09021388739347458]
           ],
           [
             [0.16139942407608032]
           ],
           [
             [0.06402886658906937]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck0_batchnorm1_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[32]
       EXLA.Backend
       [1.037217140197754, 1.2865849733352661, 0.7645959258079529, 0.6374938488006592, 1.069508671760559, 0.6837575435638428, 0.9550023674964905, 0.9839419722557068, 1.1442939043045044, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[32]
       EXLA.Backend
       [0.2217620462179184, 0.11614178121089935, 0.1586529165506363, 0.15199987590312958, 0.23966756463050842, 0.17114117741584778, 0.18509696424007416, 0.18513236939907074, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[32]
       EXLA.Backend
       [0.6259090900421143, -0.08340238034725189, 1.0731103420257568, 2.4088242053985596, -5.609257698059082, 3.4169440269470215, 7.146847724914551, ...]
     >,
     "var" => #Nx.Tensor<
       f32[32]
       EXLA.Backend
       [0.22989171743392944, 0.007035647518932819, 0.011941581964492798, 0.3442769944667816, 2.0481817722320557, 0.31869152188301086, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck6_conv1_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[192][1][3][3]
       EXLA.Backend
       [
         [
           [
             [0.014363789930939674, 0.13237811625003815, 0.09334706515073776],
             [0.12919868528842926, -0.35225632786750793, -0.1466313749551773],
             [0.06540388613939285, -0.08005677908658981, ...]
           ]
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck16_batchnorm2_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[320]
       EXLA.Backend
       [-5.189255602999765e-7, -2.7314237627251714e-7, -2.1782213366350334e-7, 1.9330661871208576e-7, 2.6874428016299134e-8, -2.4903325623881756e-8, -2.696318688322208e-7, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[320]
       EXLA.Backend
       [0.25196632742881775, 0.2486555427312851, 0.24830971658229828, 0.24681176245212555, 0.24941202998161316, 0.2485744059085846, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[320]
       EXLA.Backend
       [-0.13978023827075958, -0.05438295006752014, -0.14319643378257751, 0.1896054595708847, -0.13925714790821075, ...]
     >,
     "var" => #Nx.Tensor<
       f32[320]
       EXLA.Backend
       [0.008434775285422802, 0.007849316112697124, 0.0076168072409927845, 0.007443031296133995, ...]
     >
   },
   "mobilenetv20_features_linearbottleneck10_conv2_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[96][384][1][1]
       EXLA.Backend
       [
         [
           [
             [0.01824713684618473]
           ],
           [
             [0.08710300177335739]
           ],
           [
             [-0.10191874951124191]
           ],
           [
             [-0.07710960507392883]
           ],
           [
             [-0.042343202978372574]
           ],
           [
             [-0.014866536483168602]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck8_conv0_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[384][64][1][1]
       EXLA.Backend
       [
         [
           [
             [0.015875784680247307]
           ],
           [
             [0.00680553587153554]
           ],
           [
             [0.09295158088207245]
           ],
           [
             [-0.05411294475197792]
           ],
           [
             [0.026204092428088188]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_conv0_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[32][3][3][3]
       EXLA.Backend
       [
         [
           [
             [-0.11137320101261139, -0.22142910957336426, 0.10717468708753586],
             [4.23799006966874e-4, ...],
             ...
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck11_batchnorm2_fwd" => %{
     "beta" => #Nx.Tensor<
       f32[96]
       EXLA.Backend
       [3.9526614159512974e-7, 5.761955890193349e-7, 7.01267651948001e-7, ...]
     >,
     "gamma" => #Nx.Tensor<
       f32[96]
       EXLA.Backend
       [0.21182486414909363, 0.20167435705661774, ...]
     >,
     "mean" => #Nx.Tensor<
       f32[96]
       EXLA.Backend
       [-0.039972804486751556, ...]
     >,
     "var" => #Nx.Tensor<
       f32[96]
       EXLA.Backend
       [...]
     >
   },
   "mobilenetv20_features_linearbottleneck6_conv2_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[64][192][1][1]
       EXLA.Backend
       [
         [
           [
             [0.07365027070045471]
           ],
           [
             [-0.03652043268084526]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck3_conv2_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[32][144][1][1]
       EXLA.Backend
       [
         [
           [
             [0.02765008620917797]
           ],
           ...
         ],
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck4_conv0_fwd" => %{
     "kernel" => #Nx.Tensor<
       f32[192][32][1][1]
       EXLA.Backend
       [
         ...
       ]
     >
   },
   "mobilenetv20_features_linearbottleneck10_conv0_fwd" => %{...},
   ...
 }}
input_template = Nx.template({1, 3, target_height, target_width}, :f32)
Axon.Display.as_graph(cnn_base, input_template)
graph TD;
33["mobilenetv20_features_linearbottleneck3_conv0_fwd (:conv) {1, 144, 40, 40}"];
168["mobilenetv20_output_flatten0_reshape0 (:reshape) {1000}"];
117["mobilenetv20_features_linearbottleneck12_conv0_fwd (:conv) {1, 576, 10, 10}"];
12["mobilenetv20_features_linearbottleneck0_relu1_fwd (:relu) {1, 32, 80, 80}"];
157["mobilenetv20_features_linearbottleneck16_relu0_fwd (:relu) {1, 960, 5, 5}"];
132["mobilenetv20_features_linearbottleneck13_relu1_fwd (:relu) {1, 576, 5, 5}"];
73["mobilenetv20_features_linearbottleneck7_batchnorm1_fwd (:batch_norm) {1, 384, 20, 20}"];
44["mobilenetv20_features_linearbottleneck4_conv1_fwd (:conv) {1, 192, 20, 20}"];
124["mobilenetv20_features_linearbottleneck12_batchnorm2_fwd (:batch_norm) {1, 96, 10, 10}"];
23["mobilenetv20_features_linearbottleneck2_conv0_fwd (:conv) {1, 144, 40, 40}"];
29["mobilenetv20_features_linearbottleneck2_conv2_fwd (:conv) {1, 24, 40, 40}"];
47["mobilenetv20_features_linearbottleneck4_conv2_fwd (:conv) {1, 32, 20, 20}"];
89["mobilenetv20_features_linearbottleneck9_conv0_fwd (:conv) {1, 384, 20, 20}"];
61["mobilenetv20_features_linearbottleneck6_conv0_fwd (:conv) {1, 192, 20, 20}"];
30["mobilenetv20_features_linearbottleneck2_batchnorm2_fwd (:batch_norm) {1, 24, 40, 40}"];
43["mobilenetv20_features_linearbottleneck4_relu0_fwd (:relu) {1, 192, 20, 20}"];
163["mobilenetv20_features_conv1_fwd (:conv) {1, 1280, 5, 5}"];
39["mobilenetv20_features_linearbottleneck3_conv2_fwd (:conv) {1, 32, 20, 20}"];
131["mobilenetv20_features_linearbottleneck13_batchnorm1_fwd (:batch_norm) {1, 576, 5, 5}"];
45["mobilenetv20_features_linearbottleneck4_batchnorm1_fwd (:batch_norm) {1, 192, 20, 20}"];
48["mobilenetv20_features_linearbottleneck4_batchnorm2_fwd (:batch_norm) {1, 32, 20, 20}"];
145["mobilenetv20_features_linearbottleneck15_conv0_fwd (:conv) {1, 960, 5, 5}"];
57["mobilenetv20_features_linearbottleneck5_conv2_fwd (:conv) {1, 32, 20, 20}"];
143["container_8 (:container) {{1, 160, 5, 5}, {1, 160, 5, 5}}"];
113["mobilenetv20_features_linearbottleneck11_conv2_fwd (:conv) {1, 96, 10, 10}"];
26["mobilenetv20_features_linearbottleneck2_conv1_fwd (:conv) {1, 144, 40, 40}"];
69["mobilenetv20_features_linearbottleneck7_conv0_fwd (:conv) {1, 384, 20, 20}"];
88["add_4 (:add) {1, 64, 20, 20}"];
166["mobilenetv20_features_pool0_fwd (:global_avg_pool) {1, 1280, 1, 1}"];
144["add_8 (:add) {1, 160, 5, 5}"];
141["mobilenetv20_features_linearbottleneck14_conv2_fwd (:conv) {1, 160, 5, 5}"];
152["mobilenetv20_features_linearbottleneck15_batchnorm2_fwd (:batch_norm) {1, 160, 5, 5}"];
63["mobilenetv20_features_linearbottleneck6_relu0_fwd (:relu) {1, 192, 20, 20}"];
165["mobilenetv20_features_relu1_fwd (:relu) {1, 1280, 5, 5}"];
100["mobilenetv20_features_linearbottleneck10_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
71["mobilenetv20_features_linearbottleneck7_relu0_fwd (:relu) {1, 384, 20, 20}"];
133["mobilenetv20_features_linearbottleneck13_conv2_fwd (:conv) {1, 160, 5, 5}"];
46["mobilenetv20_features_linearbottleneck4_relu1_fwd (:relu) {1, 192, 20, 20}"];
31["container_0 (:container) {{1, 24, 40, 40}, {1, 24, 40, 40}}"];
158["mobilenetv20_features_linearbottleneck16_conv1_fwd (:conv) {1, 960, 5, 5}"];
167["mobilenetv20_output_pred_fwd (:conv) {1, 1000, 1, 1}"];
98["add_5 (:add) {1, 64, 20, 20}"];
81["mobilenetv20_features_linearbottleneck8_relu0_fwd (:relu) {1, 384, 20, 20}"];
11["mobilenetv20_features_linearbottleneck0_batchnorm1_fwd (:batch_norm) {1, 32, 80, 80}"];
37["mobilenetv20_features_linearbottleneck3_batchnorm1_fwd (:batch_norm) {1, 144, 20, 20}"];
153["container_9 (:container) {{1, 160, 5, 5}, {1, 160, 5, 5}}"];
9["mobilenetv20_features_linearbottleneck0_relu0_fwd (:relu) {1, 32, 80, 80}"];
76["mobilenetv20_features_linearbottleneck7_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
32["add_0 (:add) {1, 24, 40, 40}"];
34["mobilenetv20_features_linearbottleneck3_batchnorm0_fwd (:batch_norm) {1, 144, 40, 40}"];
25["mobilenetv20_features_linearbottleneck2_relu0_fwd (:relu) {1, 144, 40, 40}"];
28["mobilenetv20_features_linearbottleneck2_relu1_fwd (:relu) {1, 144, 40, 40}"];
85["mobilenetv20_features_linearbottleneck8_conv2_fwd (:conv) {1, 64, 20, 20}"];
91["mobilenetv20_features_linearbottleneck9_relu0_fwd (:relu) {1, 384, 20, 20}"];
120["mobilenetv20_features_linearbottleneck12_conv1_fwd (:conv) {1, 576, 10, 10}"];
64["mobilenetv20_features_linearbottleneck6_conv1_fwd (:conv) {1, 192, 20, 20}"];
146["mobilenetv20_features_linearbottleneck15_batchnorm0_fwd (:batch_norm) {1, 960, 5, 5}"];
6["mobilenetv20_features_relu0_fwd (:relu) {1, 32, 80, 80}"];
109["mobilenetv20_features_linearbottleneck11_relu0_fwd (:relu) {1, 576, 10, 10}"];
72["mobilenetv20_features_linearbottleneck7_conv1_fwd (:conv) {1, 384, 20, 20}"];
110["mobilenetv20_features_linearbottleneck11_conv1_fwd (:conv) {1, 576, 10, 10}"];
68["mobilenetv20_features_linearbottleneck6_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
86["mobilenetv20_features_linearbottleneck8_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
116["add_6 (:add) {1, 96, 10, 10}"];
51["mobilenetv20_features_linearbottleneck5_conv0_fwd (:conv) {1, 192, 20, 20}"];
80["mobilenetv20_features_linearbottleneck8_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
38["mobilenetv20_features_linearbottleneck3_relu1_fwd (:relu) {1, 144, 20, 20}"];
137["mobilenetv20_features_linearbottleneck14_relu0_fwd (:relu) {1, 960, 5, 5}"];
126["add_7 (:add) {1, 96, 10, 10}"];
13["mobilenetv20_features_linearbottleneck0_conv2_fwd (:conv) {1, 16, 80, 80}"];
59["container_2 (:container) {{1, 32, 20, 20}, {1, 32, 20, 20}}"];
40["mobilenetv20_features_linearbottleneck3_batchnorm2_fwd (:batch_norm) {1, 32, 20, 20}"];
123["mobilenetv20_features_linearbottleneck12_conv2_fwd (:conv) {1, 96, 10, 10}"];
77["container_3 (:container) {{1, 64, 20, 20}, {1, 64, 20, 20}}"];
95["mobilenetv20_features_linearbottleneck9_conv2_fwd (:conv) {1, 64, 20, 20}"];
130["mobilenetv20_features_linearbottleneck13_conv1_fwd (:conv) {1, 576, 5, 5}"];
41["mobilenetv20_features_linearbottleneck4_conv0_fwd (:conv) {1, 192, 20, 20}"];
20["mobilenetv20_features_linearbottleneck1_relu1_fwd (:relu) {1, 96, 40, 40}"];
78["add_3 (:add) {1, 64, 20, 20}"];
15["mobilenetv20_features_linearbottleneck1_conv0_fwd (:conv) {1, 96, 80, 80}"];
147["mobilenetv20_features_linearbottleneck15_relu0_fwd (:relu) {1, 960, 5, 5}"];
14["mobilenetv20_features_linearbottleneck0_batchnorm2_fwd (:batch_norm) {1, 16, 80, 80}"];
108["mobilenetv20_features_linearbottleneck11_batchnorm0_fwd (:batch_norm) {1, 576, 10, 10}"];
60["add_2 (:add) {1, 32, 20, 20}"];
162["mobilenetv20_features_linearbottleneck16_batchnorm2_fwd (:batch_norm) {1, 320, 5, 5}"];
103["mobilenetv20_features_linearbottleneck10_batchnorm1_fwd (:batch_norm) {1, 384, 10, 10}"];
7["mobilenetv20_features_linearbottleneck0_conv0_fwd (:conv) {1, 32, 80, 80}"];
122["mobilenetv20_features_linearbottleneck12_relu1_fwd (:relu) {1, 576, 10, 10}"];
140["mobilenetv20_features_linearbottleneck14_relu1_fwd (:relu) {1, 960, 5, 5}"];
83["mobilenetv20_features_linearbottleneck8_batchnorm1_fwd (:batch_norm) {1, 384, 20, 20}"];
74["mobilenetv20_features_linearbottleneck7_relu1_fwd (:relu) {1, 384, 20, 20}"];
139["mobilenetv20_features_linearbottleneck14_batchnorm1_fwd (:batch_norm) {1, 960, 5, 5}"];
99["mobilenetv20_features_linearbottleneck10_conv0_fwd (:conv) {1, 384, 20, 20}"];
112["mobilenetv20_features_linearbottleneck11_relu1_fwd (:relu) {1, 576, 10, 10}"];
97["container_5 (:container) {{1, 64, 20, 20}, {1, 64, 20, 20}}"];
149["mobilenetv20_features_linearbottleneck15_batchnorm1_fwd (:batch_norm) {1, 960, 5, 5}"];
104["mobilenetv20_features_linearbottleneck10_relu1_fwd (:relu) {1, 384, 10, 10}"];
8["mobilenetv20_features_linearbottleneck0_batchnorm0_fwd (:batch_norm) {1, 32, 80, 80}"];
134["mobilenetv20_features_linearbottleneck13_batchnorm2_fwd (:batch_norm) {1, 160, 5, 5}"];
3[/"data (:input) {1, 3, 160, 160}"/];
75["mobilenetv20_features_linearbottleneck7_conv2_fwd (:conv) {1, 64, 20, 20}"];
90["mobilenetv20_features_linearbottleneck9_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
154["add_9 (:add) {1, 160, 5, 5}"];
151["mobilenetv20_features_linearbottleneck15_conv2_fwd (:conv) {1, 160, 5, 5}"];
58["mobilenetv20_features_linearbottleneck5_batchnorm2_fwd (:batch_norm) {1, 32, 20, 20}"];
135["mobilenetv20_features_linearbottleneck14_conv0_fwd (:conv) {1, 960, 5, 5}"];
55["mobilenetv20_features_linearbottleneck5_batchnorm1_fwd (:batch_norm) {1, 192, 20, 20}"];
17["mobilenetv20_features_linearbottleneck1_relu0_fwd (:relu) {1, 96, 80, 80}"];
22["mobilenetv20_features_linearbottleneck1_batchnorm2_fwd (:batch_norm) {1, 24, 40, 40}"];
52["mobilenetv20_features_linearbottleneck5_batchnorm0_fwd (:batch_norm) {1, 192, 20, 20}"];
111["mobilenetv20_features_linearbottleneck11_batchnorm1_fwd (:batch_norm) {1, 576, 10, 10}"];
155["mobilenetv20_features_linearbottleneck16_conv0_fwd (:conv) {1, 960, 5, 5}"];
87["container_4 (:container) {{1, 64, 20, 20}, {1, 64, 20, 20}}"];
21["mobilenetv20_features_linearbottleneck1_conv2_fwd (:conv) {1, 24, 40, 40}"];
96["mobilenetv20_features_linearbottleneck9_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
4["mobilenetv20_features_conv0_fwd (:conv) {1, 32, 80, 80}"];
128["mobilenetv20_features_linearbottleneck13_batchnorm0_fwd (:batch_norm) {1, 576, 10, 10}"];
159["mobilenetv20_features_linearbottleneck16_batchnorm1_fwd (:batch_norm) {1, 960, 5, 5}"];
119["mobilenetv20_features_linearbottleneck12_relu0_fwd (:relu) {1, 576, 10, 10}"];
148["mobilenetv20_features_linearbottleneck15_conv1_fwd (:conv) {1, 960, 5, 5}"];
82["mobilenetv20_features_linearbottleneck8_conv1_fwd (:conv) {1, 384, 20, 20}"];
36["mobilenetv20_features_linearbottleneck3_conv1_fwd (:conv) {1, 144, 20, 20}"];
138["mobilenetv20_features_linearbottleneck14_conv1_fwd (:conv) {1, 960, 5, 5}"];
53["mobilenetv20_features_linearbottleneck5_relu0_fwd (:relu) {1, 192, 20, 20}"];
50["add_1 (:add) {1, 32, 20, 20}"];
66["mobilenetv20_features_linearbottleneck6_relu1_fwd (:relu) {1, 192, 20, 20}"];
24["mobilenetv20_features_linearbottleneck2_batchnorm0_fwd (:batch_norm) {1, 144, 40, 40}"];
10["mobilenetv20_features_linearbottleneck0_conv1_fwd (:conv) {1, 32, 80, 80}"];
150["mobilenetv20_features_linearbottleneck15_relu1_fwd (:relu) {1, 960, 5, 5}"];
35["mobilenetv20_features_linearbottleneck3_relu0_fwd (:relu) {1, 144, 40, 40}"];
70["mobilenetv20_features_linearbottleneck7_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
56["mobilenetv20_features_linearbottleneck5_relu1_fwd (:relu) {1, 192, 20, 20}"];
161["mobilenetv20_features_linearbottleneck16_conv2_fwd (:conv) {1, 320, 5, 5}"];
102["mobilenetv20_features_linearbottleneck10_conv1_fwd (:conv) {1, 384, 10, 10}"];
142["mobilenetv20_features_linearbottleneck14_batchnorm2_fwd (:batch_norm) {1, 160, 5, 5}"];
62["mobilenetv20_features_linearbottleneck6_batchnorm0_fwd (:batch_norm) {1, 192, 20, 20}"];
105["mobilenetv20_features_linearbottleneck10_conv2_fwd (:conv) {1, 96, 10, 10}"];
127["mobilenetv20_features_linearbottleneck13_conv0_fwd (:conv) {1, 576, 10, 10}"];
121["mobilenetv20_features_linearbottleneck12_batchnorm1_fwd (:batch_norm) {1, 576, 10, 10}"];
49["container_1 (:container) {{1, 32, 20, 20}, {1, 32, 20, 20}}"];
125["container_7 (:container) {{1, 96, 10, 10}, {1, 96, 10, 10}}"];
160["mobilenetv20_features_linearbottleneck16_relu1_fwd (:relu) {1, 960, 5, 5}"];
27["mobilenetv20_features_linearbottleneck2_batchnorm1_fwd (:batch_norm) {1, 144, 40, 40}"];
84["mobilenetv20_features_linearbottleneck8_relu1_fwd (:relu) {1, 384, 20, 20}"];
94["mobilenetv20_features_linearbottleneck9_relu1_fwd (:relu) {1, 384, 20, 20}"];
92["mobilenetv20_features_linearbottleneck9_conv1_fwd (:conv) {1, 384, 20, 20}"];
42["mobilenetv20_features_linearbottleneck4_batchnorm0_fwd (:batch_norm) {1, 192, 20, 20}"];
19["mobilenetv20_features_linearbottleneck1_batchnorm1_fwd (:batch_norm) {1, 96, 40, 40}"];
67["mobilenetv20_features_linearbottleneck6_conv2_fwd (:conv) {1, 64, 20, 20}"];
93["mobilenetv20_features_linearbottleneck9_batchnorm1_fwd (:batch_norm) {1, 384, 20, 20}"];
106["mobilenetv20_features_linearbottleneck10_batchnorm2_fwd (:batch_norm) {1, 96, 10, 10}"];
65["mobilenetv20_features_linearbottleneck6_batchnorm1_fwd (:batch_norm) {1, 192, 20, 20}"];
101["mobilenetv20_features_linearbottleneck10_relu0_fwd (:relu) {1, 384, 20, 20}"];
156["mobilenetv20_features_linearbottleneck16_batchnorm0_fwd (:batch_norm) {1, 960, 5, 5}"];
129["mobilenetv20_features_linearbottleneck13_relu0_fwd (:relu) {1, 576, 10, 10}"];
114["mobilenetv20_features_linearbottleneck11_batchnorm2_fwd (:batch_norm) {1, 96, 10, 10}"];
118["mobilenetv20_features_linearbottleneck12_batchnorm0_fwd (:batch_norm) {1, 576, 10, 10}"];
5["mobilenetv20_features_batchnorm0_fwd (:batch_norm) {1, 32, 80, 80}"];
54["mobilenetv20_features_linearbottleneck5_conv1_fwd (:conv) {1, 192, 20, 20}"];
18["mobilenetv20_features_linearbottleneck1_conv1_fwd (:conv) {1, 96, 40, 40}"];
79["mobilenetv20_features_linearbottleneck8_conv0_fwd (:conv) {1, 384, 20, 20}"];
136["mobilenetv20_features_linearbottleneck14_batchnorm0_fwd (:batch_norm) {1, 960, 5, 5}"];
115["container_6 (:container) {{1, 96, 10, 10}, {1, 96, 10, 10}}"];
107["mobilenetv20_features_linearbottleneck11_conv0_fwd (:conv) {1, 576, 10, 10}"];
164["mobilenetv20_features_batchnorm1_fwd (:batch_norm) {1, 1280, 5, 5}"];
16["mobilenetv20_features_linearbottleneck1_batchnorm0_fwd (:batch_norm) {1, 96, 80, 80}"];
167 --> 168;
166 --> 167;
165 --> 166;
164 --> 165;
163 --> 164;
162 --> 163;
161 --> 162;
160 --> 161;
159 --> 160;
158 --> 159;
157 --> 158;
156 --> 157;
155 --> 156;
154 --> 155;
153 --> 154;
144 --> 153;
152 --> 153;
151 --> 152;
150 --> 151;
149 --> 150;
148 --> 149;
147 --> 148;
146 --> 147;
145 --> 146;
144 --> 145;
143 --> 144;
134 --> 143;
142 --> 143;
141 --> 142;
140 --> 141;
139 --> 140;
138 --> 139;
137 --> 138;
136 --> 137;
135 --> 136;
134 --> 135;
133 --> 134;
132 --> 133;
131 --> 132;
130 --> 131;
129 --> 130;
128 --> 129;
127 --> 128;
126 --> 127;
125 --> 126;
116 --> 125;
124 --> 125;
123 --> 124;
122 --> 123;
121 --> 122;
120 --> 121;
119 --> 120;
118 --> 119;
117 --> 118;
116 --> 117;
115 --> 116;
106 --> 115;
114 --> 115;
113 --> 114;
112 --> 113;
111 --> 112;
110 --> 111;
109 --> 110;
108 --> 109;
107 --> 108;
106 --> 107;
105 --> 106;
104 --> 105;
103 --> 104;
102 --> 103;
101 --> 102;
100 --> 101;
99 --> 100;
98 --> 99;
97 --> 98;
88 --> 97;
96 --> 97;
95 --> 96;
94 --> 95;
93 --> 94;
92 --> 93;
91 --> 92;
90 --> 91;
89 --> 90;
88 --> 89;
87 --> 88;
78 --> 87;
86 --> 87;
85 --> 86;
84 --> 85;
83 --> 84;
82 --> 83;
81 --> 82;
80 --> 81;
79 --> 80;
78 --> 79;
77 --> 78;
68 --> 77;
76 --> 77;
75 --> 76;
74 --> 75;
73 --> 74;
72 --> 73;
71 --> 72;
70 --> 71;
69 --> 70;
68 --> 69;
67 --> 68;
66 --> 67;
65 --> 66;
64 --> 65;
63 --> 64;
62 --> 63;
61 --> 62;
60 --> 61;
59 --> 60;
50 --> 59;
58 --> 59;
57 --> 58;
56 --> 57;
55 --> 56;
54 --> 55;
53 --> 54;
52 --> 53;
51 --> 52;
50 --> 51;
49 --> 50;
40 --> 49;
48 --> 49;
47 --> 48;
46 --> 47;
45 --> 46;
44 --> 45;
43 --> 44;
42 --> 43;
41 --> 42;
40 --> 41;
39 --> 40;
38 --> 39;
37 --> 38;
36 --> 37;
35 --> 36;
34 --> 35;
33 --> 34;
32 --> 33;
31 --> 32;
22 --> 31;
30 --> 31;
29 --> 30;
28 --> 29;
27 --> 28;
26 --> 27;
25 --> 26;
24 --> 25;
23 --> 24;
22 --> 23;
21 --> 22;
20 --> 21;
19 --> 20;
18 --> 19;
17 --> 18;
16 --> 17;
15 --> 16;
14 --> 15;
13 --> 14;
12 --> 13;
11 --> 12;
10 --> 11;
9 --> 10;
8 --> 9;
7 --> 8;
6 --> 7;
5 --> 6;
4 --> 5;
3 --> 4;
{_popped, cnn_base} = cnn_base |> Axon.pop_node()
{_popped, cnn_base} = cnn_base |> Axon.pop_node()
Axon.Display.as_graph(cnn_base, input_template)
graph TD;
33["mobilenetv20_features_linearbottleneck3_conv0_fwd (:conv) {1, 144, 40, 40}"];
117["mobilenetv20_features_linearbottleneck12_conv0_fwd (:conv) {1, 576, 10, 10}"];
12["mobilenetv20_features_linearbottleneck0_relu1_fwd (:relu) {1, 32, 80, 80}"];
157["mobilenetv20_features_linearbottleneck16_relu0_fwd (:relu) {1, 960, 5, 5}"];
132["mobilenetv20_features_linearbottleneck13_relu1_fwd (:relu) {1, 576, 5, 5}"];
73["mobilenetv20_features_linearbottleneck7_batchnorm1_fwd (:batch_norm) {1, 384, 20, 20}"];
44["mobilenetv20_features_linearbottleneck4_conv1_fwd (:conv) {1, 192, 20, 20}"];
124["mobilenetv20_features_linearbottleneck12_batchnorm2_fwd (:batch_norm) {1, 96, 10, 10}"];
23["mobilenetv20_features_linearbottleneck2_conv0_fwd (:conv) {1, 144, 40, 40}"];
29["mobilenetv20_features_linearbottleneck2_conv2_fwd (:conv) {1, 24, 40, 40}"];
47["mobilenetv20_features_linearbottleneck4_conv2_fwd (:conv) {1, 32, 20, 20}"];
89["mobilenetv20_features_linearbottleneck9_conv0_fwd (:conv) {1, 384, 20, 20}"];
61["mobilenetv20_features_linearbottleneck6_conv0_fwd (:conv) {1, 192, 20, 20}"];
30["mobilenetv20_features_linearbottleneck2_batchnorm2_fwd (:batch_norm) {1, 24, 40, 40}"];
43["mobilenetv20_features_linearbottleneck4_relu0_fwd (:relu) {1, 192, 20, 20}"];
163["mobilenetv20_features_conv1_fwd (:conv) {1, 1280, 5, 5}"];
39["mobilenetv20_features_linearbottleneck3_conv2_fwd (:conv) {1, 32, 20, 20}"];
131["mobilenetv20_features_linearbottleneck13_batchnorm1_fwd (:batch_norm) {1, 576, 5, 5}"];
45["mobilenetv20_features_linearbottleneck4_batchnorm1_fwd (:batch_norm) {1, 192, 20, 20}"];
48["mobilenetv20_features_linearbottleneck4_batchnorm2_fwd (:batch_norm) {1, 32, 20, 20}"];
145["mobilenetv20_features_linearbottleneck15_conv0_fwd (:conv) {1, 960, 5, 5}"];
57["mobilenetv20_features_linearbottleneck5_conv2_fwd (:conv) {1, 32, 20, 20}"];
143["container_8 (:container) {{1, 160, 5, 5}, {1, 160, 5, 5}}"];
113["mobilenetv20_features_linearbottleneck11_conv2_fwd (:conv) {1, 96, 10, 10}"];
26["mobilenetv20_features_linearbottleneck2_conv1_fwd (:conv) {1, 144, 40, 40}"];
69["mobilenetv20_features_linearbottleneck7_conv0_fwd (:conv) {1, 384, 20, 20}"];
88["add_4 (:add) {1, 64, 20, 20}"];
166["mobilenetv20_features_pool0_fwd (:global_avg_pool) {1, 1280, 1, 1}"];
144["add_8 (:add) {1, 160, 5, 5}"];
141["mobilenetv20_features_linearbottleneck14_conv2_fwd (:conv) {1, 160, 5, 5}"];
152["mobilenetv20_features_linearbottleneck15_batchnorm2_fwd (:batch_norm) {1, 160, 5, 5}"];
63["mobilenetv20_features_linearbottleneck6_relu0_fwd (:relu) {1, 192, 20, 20}"];
165["mobilenetv20_features_relu1_fwd (:relu) {1, 1280, 5, 5}"];
100["mobilenetv20_features_linearbottleneck10_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
71["mobilenetv20_features_linearbottleneck7_relu0_fwd (:relu) {1, 384, 20, 20}"];
133["mobilenetv20_features_linearbottleneck13_conv2_fwd (:conv) {1, 160, 5, 5}"];
46["mobilenetv20_features_linearbottleneck4_relu1_fwd (:relu) {1, 192, 20, 20}"];
31["container_0 (:container) {{1, 24, 40, 40}, {1, 24, 40, 40}}"];
158["mobilenetv20_features_linearbottleneck16_conv1_fwd (:conv) {1, 960, 5, 5}"];
98["add_5 (:add) {1, 64, 20, 20}"];
81["mobilenetv20_features_linearbottleneck8_relu0_fwd (:relu) {1, 384, 20, 20}"];
11["mobilenetv20_features_linearbottleneck0_batchnorm1_fwd (:batch_norm) {1, 32, 80, 80}"];
37["mobilenetv20_features_linearbottleneck3_batchnorm1_fwd (:batch_norm) {1, 144, 20, 20}"];
153["container_9 (:container) {{1, 160, 5, 5}, {1, 160, 5, 5}}"];
9["mobilenetv20_features_linearbottleneck0_relu0_fwd (:relu) {1, 32, 80, 80}"];
76["mobilenetv20_features_linearbottleneck7_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
32["add_0 (:add) {1, 24, 40, 40}"];
34["mobilenetv20_features_linearbottleneck3_batchnorm0_fwd (:batch_norm) {1, 144, 40, 40}"];
25["mobilenetv20_features_linearbottleneck2_relu0_fwd (:relu) {1, 144, 40, 40}"];
28["mobilenetv20_features_linearbottleneck2_relu1_fwd (:relu) {1, 144, 40, 40}"];
85["mobilenetv20_features_linearbottleneck8_conv2_fwd (:conv) {1, 64, 20, 20}"];
91["mobilenetv20_features_linearbottleneck9_relu0_fwd (:relu) {1, 384, 20, 20}"];
120["mobilenetv20_features_linearbottleneck12_conv1_fwd (:conv) {1, 576, 10, 10}"];
64["mobilenetv20_features_linearbottleneck6_conv1_fwd (:conv) {1, 192, 20, 20}"];
146["mobilenetv20_features_linearbottleneck15_batchnorm0_fwd (:batch_norm) {1, 960, 5, 5}"];
6["mobilenetv20_features_relu0_fwd (:relu) {1, 32, 80, 80}"];
109["mobilenetv20_features_linearbottleneck11_relu0_fwd (:relu) {1, 576, 10, 10}"];
72["mobilenetv20_features_linearbottleneck7_conv1_fwd (:conv) {1, 384, 20, 20}"];
110["mobilenetv20_features_linearbottleneck11_conv1_fwd (:conv) {1, 576, 10, 10}"];
68["mobilenetv20_features_linearbottleneck6_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
86["mobilenetv20_features_linearbottleneck8_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
116["add_6 (:add) {1, 96, 10, 10}"];
51["mobilenetv20_features_linearbottleneck5_conv0_fwd (:conv) {1, 192, 20, 20}"];
80["mobilenetv20_features_linearbottleneck8_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
38["mobilenetv20_features_linearbottleneck3_relu1_fwd (:relu) {1, 144, 20, 20}"];
137["mobilenetv20_features_linearbottleneck14_relu0_fwd (:relu) {1, 960, 5, 5}"];
126["add_7 (:add) {1, 96, 10, 10}"];
13["mobilenetv20_features_linearbottleneck0_conv2_fwd (:conv) {1, 16, 80, 80}"];
59["container_2 (:container) {{1, 32, 20, 20}, {1, 32, 20, 20}}"];
40["mobilenetv20_features_linearbottleneck3_batchnorm2_fwd (:batch_norm) {1, 32, 20, 20}"];
123["mobilenetv20_features_linearbottleneck12_conv2_fwd (:conv) {1, 96, 10, 10}"];
77["container_3 (:container) {{1, 64, 20, 20}, {1, 64, 20, 20}}"];
95["mobilenetv20_features_linearbottleneck9_conv2_fwd (:conv) {1, 64, 20, 20}"];
130["mobilenetv20_features_linearbottleneck13_conv1_fwd (:conv) {1, 576, 5, 5}"];
41["mobilenetv20_features_linearbottleneck4_conv0_fwd (:conv) {1, 192, 20, 20}"];
20["mobilenetv20_features_linearbottleneck1_relu1_fwd (:relu) {1, 96, 40, 40}"];
78["add_3 (:add) {1, 64, 20, 20}"];
15["mobilenetv20_features_linearbottleneck1_conv0_fwd (:conv) {1, 96, 80, 80}"];
147["mobilenetv20_features_linearbottleneck15_relu0_fwd (:relu) {1, 960, 5, 5}"];
14["mobilenetv20_features_linearbottleneck0_batchnorm2_fwd (:batch_norm) {1, 16, 80, 80}"];
108["mobilenetv20_features_linearbottleneck11_batchnorm0_fwd (:batch_norm) {1, 576, 10, 10}"];
60["add_2 (:add) {1, 32, 20, 20}"];
162["mobilenetv20_features_linearbottleneck16_batchnorm2_fwd (:batch_norm) {1, 320, 5, 5}"];
103["mobilenetv20_features_linearbottleneck10_batchnorm1_fwd (:batch_norm) {1, 384, 10, 10}"];
7["mobilenetv20_features_linearbottleneck0_conv0_fwd (:conv) {1, 32, 80, 80}"];
122["mobilenetv20_features_linearbottleneck12_relu1_fwd (:relu) {1, 576, 10, 10}"];
140["mobilenetv20_features_linearbottleneck14_relu1_fwd (:relu) {1, 960, 5, 5}"];
83["mobilenetv20_features_linearbottleneck8_batchnorm1_fwd (:batch_norm) {1, 384, 20, 20}"];
74["mobilenetv20_features_linearbottleneck7_relu1_fwd (:relu) {1, 384, 20, 20}"];
139["mobilenetv20_features_linearbottleneck14_batchnorm1_fwd (:batch_norm) {1, 960, 5, 5}"];
99["mobilenetv20_features_linearbottleneck10_conv0_fwd (:conv) {1, 384, 20, 20}"];
112["mobilenetv20_features_linearbottleneck11_relu1_fwd (:relu) {1, 576, 10, 10}"];
97["container_5 (:container) {{1, 64, 20, 20}, {1, 64, 20, 20}}"];
149["mobilenetv20_features_linearbottleneck15_batchnorm1_fwd (:batch_norm) {1, 960, 5, 5}"];
104["mobilenetv20_features_linearbottleneck10_relu1_fwd (:relu) {1, 384, 10, 10}"];
8["mobilenetv20_features_linearbottleneck0_batchnorm0_fwd (:batch_norm) {1, 32, 80, 80}"];
134["mobilenetv20_features_linearbottleneck13_batchnorm2_fwd (:batch_norm) {1, 160, 5, 5}"];
3[/"data (:input) {1, 3, 160, 160}"/];
75["mobilenetv20_features_linearbottleneck7_conv2_fwd (:conv) {1, 64, 20, 20}"];
90["mobilenetv20_features_linearbottleneck9_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
154["add_9 (:add) {1, 160, 5, 5}"];
151["mobilenetv20_features_linearbottleneck15_conv2_fwd (:conv) {1, 160, 5, 5}"];
58["mobilenetv20_features_linearbottleneck5_batchnorm2_fwd (:batch_norm) {1, 32, 20, 20}"];
135["mobilenetv20_features_linearbottleneck14_conv0_fwd (:conv) {1, 960, 5, 5}"];
55["mobilenetv20_features_linearbottleneck5_batchnorm1_fwd (:batch_norm) {1, 192, 20, 20}"];
17["mobilenetv20_features_linearbottleneck1_relu0_fwd (:relu) {1, 96, 80, 80}"];
22["mobilenetv20_features_linearbottleneck1_batchnorm2_fwd (:batch_norm) {1, 24, 40, 40}"];
52["mobilenetv20_features_linearbottleneck5_batchnorm0_fwd (:batch_norm) {1, 192, 20, 20}"];
111["mobilenetv20_features_linearbottleneck11_batchnorm1_fwd (:batch_norm) {1, 576, 10, 10}"];
155["mobilenetv20_features_linearbottleneck16_conv0_fwd (:conv) {1, 960, 5, 5}"];
87["container_4 (:container) {{1, 64, 20, 20}, {1, 64, 20, 20}}"];
21["mobilenetv20_features_linearbottleneck1_conv2_fwd (:conv) {1, 24, 40, 40}"];
96["mobilenetv20_features_linearbottleneck9_batchnorm2_fwd (:batch_norm) {1, 64, 20, 20}"];
4["mobilenetv20_features_conv0_fwd (:conv) {1, 32, 80, 80}"];
128["mobilenetv20_features_linearbottleneck13_batchnorm0_fwd (:batch_norm) {1, 576, 10, 10}"];
159["mobilenetv20_features_linearbottleneck16_batchnorm1_fwd (:batch_norm) {1, 960, 5, 5}"];
119["mobilenetv20_features_linearbottleneck12_relu0_fwd (:relu) {1, 576, 10, 10}"];
148["mobilenetv20_features_linearbottleneck15_conv1_fwd (:conv) {1, 960, 5, 5}"];
82["mobilenetv20_features_linearbottleneck8_conv1_fwd (:conv) {1, 384, 20, 20}"];
36["mobilenetv20_features_linearbottleneck3_conv1_fwd (:conv) {1, 144, 20, 20}"];
138["mobilenetv20_features_linearbottleneck14_conv1_fwd (:conv) {1, 960, 5, 5}"];
53["mobilenetv20_features_linearbottleneck5_relu0_fwd (:relu) {1, 192, 20, 20}"];
50["add_1 (:add) {1, 32, 20, 20}"];
66["mobilenetv20_features_linearbottleneck6_relu1_fwd (:relu) {1, 192, 20, 20}"];
24["mobilenetv20_features_linearbottleneck2_batchnorm0_fwd (:batch_norm) {1, 144, 40, 40}"];
10["mobilenetv20_features_linearbottleneck0_conv1_fwd (:conv) {1, 32, 80, 80}"];
150["mobilenetv20_features_linearbottleneck15_relu1_fwd (:relu) {1, 960, 5, 5}"];
35["mobilenetv20_features_linearbottleneck3_relu0_fwd (:relu) {1, 144, 40, 40}"];
70["mobilenetv20_features_linearbottleneck7_batchnorm0_fwd (:batch_norm) {1, 384, 20, 20}"];
56["mobilenetv20_features_linearbottleneck5_relu1_fwd (:relu) {1, 192, 20, 20}"];
161["mobilenetv20_features_linearbottleneck16_conv2_fwd (:conv) {1, 320, 5, 5}"];
102["mobilenetv20_features_linearbottleneck10_conv1_fwd (:conv) {1, 384, 10, 10}"];
142["mobilenetv20_features_linearbottleneck14_batchnorm2_fwd (:batch_norm) {1, 160, 5, 5}"];
62["mobilenetv20_features_linearbottleneck6_batchnorm0_fwd (:batch_norm) {1, 192, 20, 20}"];
105["mobilenetv20_features_linearbottleneck10_conv2_fwd (:conv) {1, 96, 10, 10}"];
127["mobilenetv20_features_linearbottleneck13_conv0_fwd (:conv) {1, 576, 10, 10}"];
121["mobilenetv20_features_linearbottleneck12_batchnorm1_fwd (:batch_norm) {1, 576, 10, 10}"];
49["container_1 (:container) {{1, 32, 20, 20}, {1, 32, 20, 20}}"];
125["container_7 (:container) {{1, 96, 10, 10}, {1, 96, 10, 10}}"];
160["mobilenetv20_features_linearbottleneck16_relu1_fwd (:relu) {1, 960, 5, 5}"];
27["mobilenetv20_features_linearbottleneck2_batchnorm1_fwd (:batch_norm) {1, 144, 40, 40}"];
84["mobilenetv20_features_linearbottleneck8_relu1_fwd (:relu) {1, 384, 20, 20}"];
94["mobilenetv20_features_linearbottleneck9_relu1_fwd (:relu) {1, 384, 20, 20}"];
92["mobilenetv20_features_linearbottleneck9_conv1_fwd (:conv) {1, 384, 20, 20}"];
42["mobilenetv20_features_linearbottleneck4_batchnorm0_fwd (:batch_norm) {1, 192, 20, 20}"];
19["mobilenetv20_features_linearbottleneck1_batchnorm1_fwd (:batch_norm) {1, 96, 40, 40}"];
67["mobilenetv20_features_linearbottleneck6_conv2_fwd (:conv) {1, 64, 20, 20}"];
93["mobilenetv20_features_linearbottleneck9_batchnorm1_fwd (:batch_norm) {1, 384, 20, 20}"];
106["mobilenetv20_features_linearbottleneck10_batchnorm2_fwd (:batch_norm) {1, 96, 10, 10}"];
65["mobilenetv20_features_linearbottleneck6_batchnorm1_fwd (:batch_norm) {1, 192, 20, 20}"];
101["mobilenetv20_features_linearbottleneck10_relu0_fwd (:relu) {1, 384, 20, 20}"];
156["mobilenetv20_features_linearbottleneck16_batchnorm0_fwd (:batch_norm) {1, 960, 5, 5}"];
129["mobilenetv20_features_linearbottleneck13_relu0_fwd (:relu) {1, 576, 10, 10}"];
114["mobilenetv20_features_linearbottleneck11_batchnorm2_fwd (:batch_norm) {1, 96, 10, 10}"];
118["mobilenetv20_features_linearbottleneck12_batchnorm0_fwd (:batch_norm) {1, 576, 10, 10}"];
5["mobilenetv20_features_batchnorm0_fwd (:batch_norm) {1, 32, 80, 80}"];
54["mobilenetv20_features_linearbottleneck5_conv1_fwd (:conv) {1, 192, 20, 20}"];
18["mobilenetv20_features_linearbottleneck1_conv1_fwd (:conv) {1, 96, 40, 40}"];
79["mobilenetv20_features_linearbottleneck8_conv0_fwd (:conv) {1, 384, 20, 20}"];
136["mobilenetv20_features_linearbottleneck14_batchnorm0_fwd (:batch_norm) {1, 960, 5, 5}"];
115["container_6 (:container) {{1, 96, 10, 10}, {1, 96, 10, 10}}"];
107["mobilenetv20_features_linearbottleneck11_conv0_fwd (:conv) {1, 576, 10, 10}"];
164["mobilenetv20_features_batchnorm1_fwd (:batch_norm) {1, 1280, 5, 5}"];
16["mobilenetv20_features_linearbottleneck1_batchnorm0_fwd (:batch_norm) {1, 96, 80, 80}"];
165 --> 166;
164 --> 165;
163 --> 164;
162 --> 163;
161 --> 162;
160 --> 161;
159 --> 160;
158 --> 159;
157 --> 158;
156 --> 157;
155 --> 156;
154 --> 155;
153 --> 154;
144 --> 153;
152 --> 153;
151 --> 152;
150 --> 151;
149 --> 150;
148 --> 149;
147 --> 148;
146 --> 147;
145 --> 146;
144 --> 145;
143 --> 144;
134 --> 143;
142 --> 143;
141 --> 142;
140 --> 141;
139 --> 140;
138 --> 139;
137 --> 138;
136 --> 137;
135 --> 136;
134 --> 135;
133 --> 134;
132 --> 133;
131 --> 132;
130 --> 131;
129 --> 130;
128 --> 129;
127 --> 128;
126 --> 127;
125 --> 126;
116 --> 125;
124 --> 125;
123 --> 124;
122 --> 123;
121 --> 122;
120 --> 121;
119 --> 120;
118 --> 119;
117 --> 118;
116 --> 117;
115 --> 116;
106 --> 115;
114 --> 115;
113 --> 114;
112 --> 113;
111 --> 112;
110 --> 111;
109 --> 110;
108 --> 109;
107 --> 108;
106 --> 107;
105 --> 106;
104 --> 105;
103 --> 104;
102 --> 103;
101 --> 102;
100 --> 101;
99 --> 100;
98 --> 99;
97 --> 98;
88 --> 97;
96 --> 97;
95 --> 96;
94 --> 95;
93 --> 94;
92 --> 93;
91 --> 92;
90 --> 91;
89 --> 90;
88 --> 89;
87 --> 88;
78 --> 87;
86 --> 87;
85 --> 86;
84 --> 85;
83 --> 84;
82 --> 83;
81 --> 82;
80 --> 81;
79 --> 80;
78 --> 79;
77 --> 78;
68 --> 77;
76 --> 77;
75 --> 76;
74 --> 75;
73 --> 74;
72 --> 73;
71 --> 72;
70 --> 71;
69 --> 70;
68 --> 69;
67 --> 68;
66 --> 67;
65 --> 66;
64 --> 65;
63 --> 64;
62 --> 63;
61 --> 62;
60 --> 61;
59 --> 60;
50 --> 59;
58 --> 59;
57 --> 58;
56 --> 57;
55 --> 56;
54 --> 55;
53 --> 54;
52 --> 53;
51 --> 52;
50 --> 51;
49 --> 50;
40 --> 49;
48 --> 49;
47 --> 48;
46 --> 47;
45 --> 46;
44 --> 45;
43 --> 44;
42 --> 43;
41 --> 42;
40 --> 41;
39 --> 40;
38 --> 39;
37 --> 38;
36 --> 37;
35 --> 36;
34 --> 35;
33 --> 34;
32 --> 33;
31 --> 32;
22 --> 31;
30 --> 31;
29 --> 30;
28 --> 29;
27 --> 28;
26 --> 27;
25 --> 26;
24 --> 25;
23 --> 24;
22 --> 23;
21 --> 22;
20 --> 21;
19 --> 20;
18 --> 19;
17 --> 18;
16 --> 17;
15 --> 16;
14 --> 15;
13 --> 14;
12 --> 13;
11 --> 12;
10 --> 11;
9 --> 10;
8 --> 9;
7 --> 8;
6 --> 7;
5 --> 6;
4 --> 5;
3 --> 4;

Building the Model from Mobilenet Model

model =
  cnn_base
  |> Axon.namespace("feature_extractor")
  |> Axon.freeze()
  |> Axon.global_avg_pool(channels: :first)
  |> Axon.dropout(rate: 0.2)
  |> Axon.dense(1)
#Axon<
  inputs: %{"data" => {1, 3, 224, 224}}
  outputs: "dense_0"
  nodes: 168
>

Training the Model

loss =
  &amp;Axon.Losses.binary_cross_entropy(&amp;1, &amp;2,
    reduction: :mean,
    from_logits: true
  )

optimizer = Polaris.Optimizers.adam(learning_rate: 1.0e-4)

trained_model_state =
  model
  |> Axon.Loop.trainer(loss, optimizer)
  |> Axon.Loop.metric(:accuracy)
  |> Axon.Loop.validate(model, val_pipeline)
  |> Axon.Loop.early_stop("validation_loss", mode: :min, patience: 5)
  |> Axon.Loop.run(
    train_pipeline,
    %{"feature_extractor" => cnn_base_params},
    epochs: 100,
    compiler: EXLA
  )

18:00:45.817 [debug] Forwarding options: [compiler: EXLA] to JIT compiler

18:00:45.952 [warning] found unexpected key in the initial parameters map: "mobilenetv20_output_pred_fwd"
Epoch: 0, Batch: 700, accuracy: 0.7490190 loss: 0.4872290
Epoch: 1, Batch: 700, accuracy: 0.8607792 loss: 0.3941975
Epoch: 2, Batch: 700, accuracy: 0.8781650 loss: 0.3532889
Epoch: 3, Batch: 700, accuracy: 0.8869918 loss: 0.3276705
Epoch: 4, Batch: 700, accuracy: 0.8916724 loss: 0.3102300
Epoch: 5, Batch: 700, accuracy: 0.8947041 loss: 0.2979036
Epoch: 6, Batch: 700, accuracy: 0.8957738 loss: 0.2888097
Epoch: 7, Batch: 700, accuracy: 0.9012572 loss: 0.2812564
Epoch: 8, Batch: 700, accuracy: 0.8998307 loss: 0.2750721
Epoch: 9, Batch: 700, accuracy: 0.9046897 loss: 0.2699555
Epoch: 10, Batch: 700, accuracy: 0.9046898 loss: 0.2652446
Epoch: 11, Batch: 700, accuracy: 0.9021933 loss: 0.2615432
Epoch: 12, Batch: 700, accuracy: 0.9069632 loss: 0.2578678
Epoch: 13, Batch: 700, accuracy: 0.9043331 loss: 0.2548892
Epoch: 14, Batch: 700, accuracy: 0.9046451 loss: 0.2521417
Epoch: 15, Batch: 700, accuracy: 0.9068741 loss: 0.2496624
Epoch: 16, Batch: 700, accuracy: 0.9083008 loss: 0.2473769
Epoch: 17, Batch: 700, accuracy: 0.9061162 loss: 0.2456258
Epoch: 18, Batch: 700, accuracy: 0.9071863 loss: 0.2438199
Epoch: 19, Batch: 700, accuracy: 0.9093258 loss: 0.2419264
Epoch: 20, Batch: 700, accuracy: 0.9084345 loss: 0.2402944
Epoch: 21, Batch: 700, accuracy: 0.9063838 loss: 0.2390028
Epoch: 22, Batch: 700, accuracy: 0.9085682 loss: 0.2377930
Epoch: 23, Batch: 700, accuracy: 0.9099947 loss: 0.2364875
Epoch: 24, Batch: 700, accuracy: 0.9115102 loss: 0.2352752
Epoch: 25, Batch: 700, accuracy: 0.9088356 loss: 0.2342139
Epoch: 26, Batch: 700, accuracy: 0.9062501 loss: 0.2333108
Epoch: 27, Batch: 700, accuracy: 0.9084344 loss: 0.2323771
Epoch: 28, Batch: 700, accuracy: 0.9125357 loss: 0.2313402
Epoch: 29, Batch: 700, accuracy: 0.9106188 loss: 0.2304147
Epoch: 30, Batch: 700, accuracy: 0.9093260 loss: 0.2297038
Epoch: 31, Batch: 700, accuracy: 0.9107079 loss: 0.2288840
Epoch: 32, Batch: 700, accuracy: 0.9102621 loss: 0.2282075
Epoch: 33, Batch: 700, accuracy: 0.9108863 loss: 0.2274254
Epoch: 34, Batch: 700, accuracy: 0.9124019 loss: 0.2267024
Epoch: 35, Batch: 700, accuracy: 0.9111538 loss: 0.2260765
Epoch: 36, Batch: 700, accuracy: 0.9125357 loss: 0.2254625
Epoch: 37, Batch: 700, accuracy: 0.9126249 loss: 0.2248283
Epoch: 38, Batch: 700, accuracy: 0.9089692 loss: 0.2243719
Epoch: 39, Batch: 700, accuracy: 0.9092367 loss: 0.2239486
Epoch: 40, Batch: 700, accuracy: 0.9117332 loss: 0.2234031
Epoch: 41, Batch: 700, accuracy: 0.9137395 loss: 0.2228346
Epoch: 42, Batch: 700, accuracy: 0.9102621 loss: 0.2224314
Batch: 6, accuracy: 0.9464286 loss: 0.1156730
...
        EXLA.Backend
        [0.1144828125834465, 0.29103899002075195, 0.19599959254264832, 0.11088918149471283, 0.11658263951539993, 0.2404186874628067, 0.05210668221116066, 0.279130220413208, 0.16053666174411774, 0.10330193489789963, 0.17037393152713776, 0.42853739857673645, 0.2756047546863556, 0.3480927646160126, 0.41298505663871765, 0.3465864956378937, 0.2918163239955902, 0.2078855335712433, 0.19570542871952057, 0.3232663571834564, 0.20460295677185059, 0.21620818972587585, 0.06636220216751099, 0.16237126290798187, 0.11230707168579102, 0.20659472048282623, 0.04134644940495491, 0.16707743704319, 0.18970169126987457, 0.08170202374458313, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [7.723511662334204e-4, -3.091116959694773e-4, -9.241990483133122e-5, -6.965994543861598e-5, 2.6191226424998604e-5, -3.1657828367315233e-4, 1.1303959036013111e-4, 1.4945304428692907e-4, 4.511853330768645e-4, -4.4339807936921716e-4, -1.6256280650850385e-4, 1.3352044334169477e-4, 7.505110261263326e-5, 5.360154318623245e-4, 0.001765010878443718, -3.995028673671186e-4, -1.3021344784647226e-5, 7.265885942615569e-4, 4.0445878403261304e-4, -2.8110373023082502e-5, 1.516808697488159e-4, 1.1108114995295182e-4, -1.617355301277712e-4, 4.557485517580062e-4, 1.0374527482781559e-4, -6.70250374241732e-5, 1.5719550719950348e-4, 3.4182274248450994e-5, -4.4115696800872684e-4, ...]
      >,
      "var" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [0.09834592044353485, 0.1735430806875229, 0.38497281074523926, 0.17505548894405365, 0.19528056681156158, 0.1011020764708519, 0.18303343653678894, 0.20173640549182892, 0.09403466433286667, 0.06812598556280136, 0.14213919639587402, 0.43439722061157227, 0.1752435863018036, 0.4183683693408966, 0.3139973282814026, 0.3875087797641754, 0.09094872325658798, 0.10979244112968445, 0.3352102041244507, 0.11777494847774506, 0.193935826420784, 0.1976424902677536, 0.1345539391040802, 0.2652271091938019, 0.16767887771129608, 0.1588459610939026, 0.09324903041124344, 0.2787035405635834, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck11_batchnorm0_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [0.1379205733537674, 0.1852763593196869, 0.07007614523172379, 0.14786551892757416, 0.08499550819396973, 0.022980524227023125, 0.10445888340473175, -0.03380206599831581, 0.05621454492211342, 0.13879631459712982, 0.1278844028711319, 0.11503273248672485, 0.008648642338812351, 0.07986151427030563, 0.05827014893293381, 0.02635921724140644, 0.16676078736782074, 0.1696624606847763, -0.0049711973406374454, 0.11680029332637787, 0.09153503179550171, -0.09545790404081345, 0.10052905976772308, 0.2198638916015625, 0.09066508710384369, 0.021357741206884384, 0.1243291050195694, 0.13527752459049225, 0.09010029584169388, 0.07285095006227493, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [0.1116461232304573, 0.08017633110284805, 0.11955763399600983, 0.09390335530042648, 0.09855279326438904, 0.14738331735134125, 0.0888434574007988, 0.13319958746433258, 0.10988157987594604, 0.09758485853672028, 0.10826665163040161, 0.1016610860824585, 0.09888888150453568, 0.12156049907207489, 0.11694491654634476, 0.119488425552845, 0.07615012675523758, 0.06269187480211258, 0.15567438304424286, 0.1064029410481453, 0.11337470263242722, 0.14443783462047577, 0.0840112492442131, 0.09276950359344482, 0.16790293157100677, 0.11844969540834427, 0.0930234044790268, 0.11310834437608719, 0.1162751242518425, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [3.63793503765919e-7, 1.209602089602413e-7, 5.4800288751266635e-8, 2.552324076532386e-7, 3.2612440747925575e-8, 1.4689121030642127e-8, 1.1114344999896275e-7, -8.35849434110969e-8, 2.2601201976613083e-7, 1.6091770760340296e-7, 3.1186812066152925e-7, -1.1275391287313141e-8, -2.748528871165945e-8, -1.7028881416081276e-7, -1.3741856719207135e-7, -1.3034306789450056e-7, 4.320734205975896e-7, 5.223997732173302e-7, 1.3161741918565895e-8, 1.749855726984606e-7, 6.074458980265263e-8, 4.904541128780693e-7, 2.8845341670802327e-9, 3.752578692228781e-8, -8.265797646345163e-8, -2.6385482243540537e-8, 6.387549156272598e-9, -3.501408230022207e-8, ...]
      >,
      "var" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [0.02586364559829235, 0.04419300705194473, 0.02695547789335251, 0.027057670056819916, 0.023040946573019028, 0.03236440196633339, 0.021965283900499344, 0.03861037641763687, 0.02072414942085743, 0.020027386024594307, 0.017105598002672195, 0.027752427384257317, 0.025669051334261894, 0.030634647235274315, 0.02309420332312584, 0.026124022901058197, 0.009496251121163368, 0.02531951665878296, 0.04465075209736824, 0.03168340027332306, 0.03208788484334946, 0.03684369474649429, 0.01707514189183712, 0.033794183284044266, 0.05672197788953781, 0.01841931976377964, 0.031393010169267654, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck15_batchnorm1_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [-0.08277615904808044, 0.12612955272197723, -0.007448040880262852, -0.04266936331987381, 0.07737076282501221, 0.0038264743052423, -0.04381075128912926, -0.061556603759527206, 0.07451092451810837, 0.00787589605897665, -0.020372798666357994, 0.03865779936313629, -0.05657007917761803, -0.048030927777290344, -0.07944952696561813, 0.07044878602027893, 0.009796769358217716, -0.0038760113529860973, -0.06943286210298538, 0.035002585500478745, -0.06959164142608643, -0.08562169969081879, -0.1489018350839615, -0.0025909075047820807, -0.01582941599190235, -0.0698612853884697, 0.05639493837952614, 0.3489030599594116, -0.032863155007362366, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [0.3150080740451813, 0.16160078346729279, 0.22753824293613434, 0.23682628571987152, 0.16329334676265717, 0.1377934366464615, 0.19716574251651764, 0.15468646585941315, 0.14866210520267487, 0.1379021555185318, 0.18138086795806885, 0.1989935338497162, 0.2044813334941864, 0.22123073041439056, 0.22492459416389465, 0.158174529671669, 0.14002622663974762, 0.13364093005657196, 0.19789022207260132, 0.17698919773101807, 0.1932896226644516, 0.17972171306610107, 0.24500404298305511, 0.1422407180070877, 0.2906820476055145, 0.21549536287784576, 0.16308309137821198, 0.116135373711586, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [0.006607928778976202, -0.005489333998411894, 0.001202044659294188, 0.016741983592510223, -0.021014627069234848, 3.9149439544416964e-4, -0.001380614354275167, 0.006143252365291119, -0.022377200424671173, 3.921858442481607e-4, 0.011808911338448524, -0.038100697100162506, -0.0027107521891593933, -0.002684985054656863, -0.003361391369253397, -0.027487538754940033, 0.0032275188714265823, 7.377289002761245e-4, -0.003101653652265668, -0.02393931709229946, -0.0026154580991715193, 4.0451946551911533e-4, 0.019924011081457138, 0.0014483053237199783, 0.012146278284490108, -6.006719195283949e-4, -0.02536875568330288, ...]
      >,
      "var" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [6.60044839605689e-4, 2.093793300446123e-4, 6.339335814118385e-4, 8.293294231407344e-4, 4.367101937532425e-4, 3.0174001949490048e-5, 1.974338520085439e-4, 5.38308268005494e-5, 5.038247327320278e-4, 1.9694401999004185e-4, 3.8651813520118594e-4, 5.852532922290266e-4, 2.917999809142202e-4, 3.8506241980940104e-4, 3.85848106816411e-4, 4.808907106053084e-4, 8.100475679384544e-5, 6.1370710682240315e-6, 1.8867437029257417e-4, 3.83470905944705e-4, 5.22851652931422e-4, 4.0124301449395716e-5, 4.517717461567372e-4, 7.429129254887812e-6, 0.0013210318284109235, 2.3677274293731898e-4, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck9_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[64][384][1][1]
        EXLA.Backend
        [
          [
            [
              [-0.029621589928865433]
            ],
            [
              [0.03196175768971443]
            ],
            [
              [-0.04775437340140343]
            ],
            [
              [-0.02047593705356121]
            ],
            [
              [0.022681893780827522]
            ],
            [
              [0.0061656758189201355]
            ],
            [
              [-0.020124075934290886]
            ],
            [
              [-0.0419141948223114]
            ],
            [
              [-0.0958712100982666]
            ],
            [
              [-0.03112773969769478]
            ],
            [
              [0.033005211502313614]
            ],
            [
              [-0.0796334370970726]
            ],
            [
              [-0.08826649934053421]
            ],
            [
              [0.008132039569318295]
            ],
            [
              [0.13005898892879486]
            ],
            [
              [0.01112390961498022]
            ],
            [
              [-0.025056127458810806]
            ],
            [
              [-0.021656425669789314]
            ],
            [
              [-0.009619093500077724]
            ],
            [
              [-0.09488093852996826]
            ],
            [
              [-0.03190577030181885]
            ],
            [
              [0.08531086146831512]
            ],
            [
              [0.047572940587997437]
            ],
            [
              [-0.036207523196935654]
            ],
            [
              [-0.1315651535987854]
            ],
            [
              [-0.03746785223484039]
            ],
            [
              [-0.04043968394398689]
            ],
            [
              [-0.1008371040225029]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck11_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[96][576][1][1]
        EXLA.Backend
        [
          [
            [
              [-0.04216459393501282]
            ],
            [
              [-0.03359765186905861]
            ],
            [
              [0.021302534267306328]
            ],
            [
              [0.06966963410377502]
            ],
            [
              [-0.023086098954081535]
            ],
            [
              [0.02423691935837269]
            ],
            [
              [-0.0027191585395485163]
            ],
            [
              [-0.12812955677509308]
            ],
            [
              [0.10142464935779572]
            ],
            [
              [0.04367773234844208]
            ],
            [
              [-0.04808792844414711]
            ],
            [
              [-0.059124793857336044]
            ],
            [
              [-0.03544275462627411]
            ],
            [
              [-0.10292188823223114]
            ],
            [
              [-0.15540222823619843]
            ],
            [
              [0.08993023633956909]
            ],
            [
              [-0.0018712839810177684]
            ],
            [
              [0.04499407112598419]
            ],
            [
              [0.03717062994837761]
            ],
            [
              [0.1088862419128418]
            ],
            [
              [0.03789613023400307]
            ],
            [
              [0.011467562057077885]
            ],
            [
              [0.022153226658701897]
            ],
            [
              [0.011087901890277863]
            ],
            [
              [0.10830669850111008]
            ],
            [
              [-0.039926622062921524]
            ],
            [
              [0.03126271069049835]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck2_batchnorm1_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[144]
        EXLA.Backend
        [0.06266669183969498, 0.34174367785453796, 0.014498202130198479, 0.09171395748853683, -0.01248596515506506, 0.09248802065849304, 0.03913332521915436, -0.12479757517576218, 0.04716098681092262, 0.2623796761035919, 0.00875355675816536, 0.3866884410381317, 0.34793952107429504, 0.07759835571050644, 0.28590700030326843, 0.34494516253471375, 0.1602264642715454, 0.3233429491519928, -0.08500878512859344, 0.017791876569390297, 0.012547487393021584, 0.07443935424089432, 0.470037043094635, 0.06998748332262039, -0.1145705133676529, 0.35432684421539307, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[144]
        EXLA.Backend
        [0.14569389820098877, 0.08483350276947021, 0.18709848821163177, 0.08500487357378006, 0.0695008710026741, 0.13266052305698395, 0.23066304624080658, 0.1294335573911667, 0.15951289236545563, 0.1612233966588974, 0.04986649379134178, 0.24220621585845947, 0.13059896230697632, 0.2140040248632431, 0.16401590406894684, 0.04375945031642914, 0.1284344643354416, 0.10862768441438675, 0.22237107157707214, 0.20400500297546387, 0.2341395914554596, 0.36944401264190674, 0.1626269817352295, 0.1425088793039322, 0.37626516819000244, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[144]
        EXLA.Backend
        [0.04842672124505043, 0.011466434225440025, -0.28272300958633423, 0.04086904227733612, -0.0, 0.03728378936648369, -0.07810080051422119, -0.0095414062961936, 0.019914837554097176, -0.26291584968566895, 1.200515914815965e-22, 0.4378272294998169, -0.07176417112350464, -0.20427647233009338, -0.40337827801704407, 0.07497201859951019, 0.06511131674051285, 0.02972135879099369, -0.029825769364833832, 0.03816595301032066, 0.008758057840168476, -0.5172790288925171, 0.18057824671268463, 0.03219117224216461, ...]
      >,
      "var" => #Nx.Tensor<
        f32[144]
        EXLA.Backend
        [0.0016082772053778172, 0.008874273858964443, 0.014802338555455208, 0.010124930180609226, 0.0, 0.0027794779743999243, 0.019753890112042427, 6.709994049742818e-4, 0.004814377520233393, 0.018505984917283058, 3.875728477835647e-26, 0.024166392162442207, 0.011687243357300758, 0.018400687724351883, 0.00805666297674179, 0.003925488330423832, 0.004267431795597076, 0.005561113357543945, 0.00586837250739336, 0.01074590440839529, 0.023998573422431946, 0.010233975946903229, 0.007847330532968044, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck3_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[144][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.007945280522108078, -0.5135775208473206, 0.642527163028717],
              [0.09571424126625061, -0.4191504120826721, -0.42046400904655457],
              [0.2738749086856842, -1.341597080230713, 1.7941224575042725]
            ]
          ],
          [
            [
              [0.36698344349861145, 0.8006415963172913, 1.0355921983718872],
              [0.2570298910140991, 0.9058852791786194, 0.1435706615447998],
              [0.08509054780006409, -0.29565608501434326, -0.19910597801208496]
            ]
          ],
          [
            [
              [-0.07826036959886551, 1.5095280408859253, -0.7258926033973694],
              [-0.1067059189081192, -0.07988816499710083, -0.4704415500164032],
              [-2.0861942768096924, ...]
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck10_batchnorm0_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [-0.14200253784656525, -0.08138399571180344, 0.020918697118759155, -0.030422303825616837, -0.054856158792972565, 0.0021326385904103518, -0.18490885198116302, -0.03978404402732849, 0.17561353743076324, 0.007802810054272413, -0.05954292416572571, 0.001685986528173089, -0.021205948665738106, 0.09003060311079025, 0.03726910054683685, 0.14319199323654175, 0.15906931459903717, 0.09075694531202316, 0.14122611284255981, 0.19317775964736938, 0.11409415304660797, 0.2862238883972168, 0.21863330900669098, 0.31881415843963623, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [0.08955207467079163, 0.12251311540603638, 0.13474930822849274, 0.12169622629880905, 0.10408146679401398, 0.12673817574977875, 0.09429759532213211, 0.004978157579898834, 0.06450905650854111, 0.12292398512363434, 0.008911419659852982, 0.16245806217193604, 0.12929747998714447, 0.07689117640256882, 0.13534259796142578, 0.0939403846859932, 0.07975732535123825, 0.07448924332857132, 0.02857014909386635, 0.05634821206331253, 0.04961058497428894, 0.08511262387037277, 0.1229783222079277, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [2.3069126473274082e-4, 8.844963740557432e-4, 2.083192957798019e-4, 2.62250512605533e-4, -1.6108462295960635e-4, 3.3244339283555746e-4, 1.4792397269047797e-4, 8.78396604093723e-5, -5.305111990310252e-4, 3.3695020829327404e-4, 1.6378094733227044e-4, 8.148422930389643e-4, 4.689058114308864e-4, 2.057998499367386e-4, 7.603167905472219e-4, -1.234020892297849e-4, -2.2832311515230685e-4, 1.4915215433575213e-4, -1.732776036078576e-5, 1.1797493789345026e-4, 3.793840587604791e-4, 8.259260539489333e-6, ...]
      >,
      "var" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [0.0380789078772068, 0.19644825160503387, 0.16667498648166656, 0.12124160677194595, 0.0748945027589798, 0.1588868945837021, 0.06828655302524567, 0.003897328395396471, 0.3409107029438019, 0.15085452795028687, 0.011681756936013699, 0.3451366722583771, 0.19449232518672943, 0.12895861268043518, 0.2786988615989685, 0.1285005509853363, 0.13332460820674896, 0.09845466166734695, 0.12204387784004211, 0.20885948836803436, 0.12294120341539383, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck14_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[160]
        EXLA.Backend
        [-8.388164474126825e-8, -5.7530677111117257e-8, 7.027283999150313e-8, -9.3144599588868e-8, 4.638787132194011e-8, -9.98987204070545e-8, 1.4243799739688257e-7, -1.1660619492204205e-7, 6.76579361424956e-8, 1.2566950147174794e-7, -1.1262851984383815e-7, 6.708535948973804e-8, 2.0174957171548158e-7, 1.1858566040245933e-7, 3.9078233982081656e-8, -1.565037024420235e-7, 1.173001891174863e-7, -1.3883133931358316e-7, -1.3765874484761298e-7, -1.1415199452358138e-7, -3.374253765286994e-8, 1.203010810968408e-7, -1.509711751168652e-7, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[160]
        EXLA.Backend
        [0.18752147257328033, 0.34522393345832825, 0.24475780129432678, 0.33060014247894287, 0.21970151364803314, 0.3153970539569855, 0.28512927889823914, 0.169156014919281, 0.3166297674179077, 0.2279355376958847, 0.21180318295955658, 0.22759808599948883, 0.23080326616764069, 0.17287319898605347, 0.27370786666870117, 0.2255641222000122, 0.2651505172252655, 0.17854700982570648, 0.257531076669693, 0.2266509234905243, 0.2487022578716278, 0.22212600708007812, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[160]
        EXLA.Backend
        [-0.10289469361305237, 0.13421736657619476, 0.3127436637878418, -0.010098671540617943, 0.11501706391572952, 0.07170626521110535, -0.12338873744010925, -0.09662780910730362, 0.10249986499547958, 0.039654918015003204, 0.23144079744815826, -0.02762533910572529, -0.2323276251554489, -0.14802084863185883, -0.17405559122562408, -0.05928152799606323, 0.30801624059677124, 0.03448152542114258, 0.10294467210769653, -0.14233750104904175, 0.042918454855680466, ...]
      >,
      "var" => #Nx.Tensor<
        f32[160]
        EXLA.Backend
        [0.01982937380671501, 0.07223469018936157, 0.03559407591819763, 0.05254514887928963, 0.03399147093296051, 0.05125325173139572, 0.045940618962049484, 0.015112430788576603, 0.0665295273065567, 0.029054192826151848, 0.024110015481710434, 0.028089039027690887, 0.032291218638420105, 0.017801793292164803, 0.04934801161289215, 0.028767650946974754, 0.036390747874975204, 0.020782455801963806, 0.038260187953710556, 0.03192899003624916, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck0_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[16]
        EXLA.Backend
        [-0.0010560675291344523, -6.787155289202929e-4, -7.046025712043047e-4, -9.398604743182659e-4, -8.966192253865302e-4, -2.11737904464826e-4, -5.149749922566116e-4, -0.0010324727045372128, -7.925425306893885e-4, -9.090156527236104e-5, -4.453391011338681e-4, -9.805577574297786e-4, -7.45584664400667e-4, -7.640189724043012e-4, -0.0011470100143924356, -3.8889769348315895e-4]
      >,
      "gamma" => #Nx.Tensor<
        f32[16]
        EXLA.Backend
        [0.5097491145133972, 0.6192775368690491, 0.657399594783783, 0.5558041930198669, 0.6164559721946716, 0.6381960511207581, 0.6152222156524658, 0.4946329891681671, 0.5616724491119385, 0.6567564606666565, 0.5526584982872009, 0.42007970809936523, 0.4637764096260071, 0.5626768469810486, 0.6189385056495667, 0.6354565620422363]
      >,
      "mean" => #Nx.Tensor<
        f32[16]
        EXLA.Backend
        [0.021046694368124008, -0.4771552085876465, -0.4624612629413605, 1.7339810132980347, 0.3934212327003479, -0.845706045627594, 1.2445530891418457, -0.4817613959312439, 0.8393087983131409, -2.1945111751556396, 2.2859015464782715, 0.042962949723005295, -1.1515218019485474, -0.37790021300315857, -3.0484421253204346, -0.784759521484375]
      >,
      "var" => #Nx.Tensor<
        f32[16]
        EXLA.Backend
        [0.015569682233035564, 0.057801131159067154, 0.05658106878399849, 0.018852993845939636, 0.038473501801490784, 0.0927843227982521, 0.05347060412168503, 0.04166773706674576, 0.037543755024671555, 0.14158287644386292, 0.04085716977715492, 0.060751400887966156, 0.034666694700717926, 0.056090183556079865, 0.04101115092635155, 0.052235011011362076]
      >
    },
    "mobilenetv20_features_linearbottleneck7_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[384][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.09069628268480301, 0.011556080542504787, -0.044405367225408554],
              [0.1974998563528061, -0.05144780874252319, -0.15406860411167145],
              [0.03294772282242775, -0.08055609464645386, -0.10401368886232376]
            ]
          ],
          [
            [
              [0.004530614707618952, -0.016135336831212044, 0.010837683454155922],
              [-0.013911036774516106, -0.17941752076148987, 0.0742872878909111],
              [0.03701423481106758, 0.17489252984523773, 0.049539677798748016]
            ]
          ],
          [
            [
              [-0.03442509472370148, -0.04338786378502846, -8.070958429016173e-4],
              ...
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck1_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[96][16][1][1]
        EXLA.Backend
        [
          [
            [
              [-0.26627230644226074]
            ],
            [
              [0.1223592683672905]
            ],
            [
              [-0.17675793170928955]
            ],
            [
              [-0.14175795018672943]
            ],
            [
              [-0.04847275838255882]
            ],
            [
              [0.02572016417980194]
            ],
            [
              [0.18247880041599274]
            ],
            [
              [-0.22272470593452454]
            ],
            [
              [-0.0375950001180172]
            ],
            [
              [-0.12036115676164627]
            ],
            [
              [-0.42806223034858704]
            ],
            [
              [0.2049194574356079]
            ],
            [
              [0.025416148826479912]
            ],
            [
              [-0.20868347585201263]
            ],
            [
              [-0.00869692862033844]
            ],
            [
              [0.10014521330595016]
            ]
          ],
          [
            [
              [0.1901395469903946]
            ],
            [
              [-0.15854698419570923]
            ],
            [
              [-0.05858785659074783]
            ],
            [
              [-0.016011832281947136]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck2_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[24]
        EXLA.Backend
        [-6.408389890566468e-5, -1.0128262510988861e-4, -4.3479518353706226e-5, -1.589354214956984e-4, -7.557999197160825e-5, -2.3196584152174182e-5, -2.216532493548584e-6, 9.525781933916733e-6, -8.600373985245824e-5, -2.3008385323919356e-4, -2.3174902889877558e-4, -2.012878976529464e-4, -9.838923870120198e-5, -1.22729703434743e-4, -9.646145917940885e-5, -2.797471643134486e-5, -1.039204653352499e-4, -1.2249314750079066e-4, -8.16335596027784e-5, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[24]
        EXLA.Backend
        [0.5065965056419373, 0.44786974787712097, 0.5534567832946777, 0.3781253695487976, 0.44674113392829895, 0.28593751788139343, 0.5029821395874023, 0.19650447368621826, 0.3918337821960449, 0.33050450682640076, 0.31709131598472595, 0.1402808576822281, 0.4198921024799347, 0.02611447498202324, 0.5337468981742859, 0.49816271662712097, 0.32386210560798645, 0.40961867570877075, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[24]
        EXLA.Backend
        [-0.49676281213760376, -0.5412184000015259, 0.2777172029018402, 0.30402812361717224, -0.038367077708244324, 0.19269412755966187, -0.20861922204494476, 0.3330950438976288, -0.35258007049560547, -0.0413212850689888, -0.44278526306152344, 0.43556883931159973, -0.007557356264442205, -0.049789782613515854, -0.1810888946056366, -0.1666330248117447, 0.15477608144283295, ...]
      >,
      "var" => #Nx.Tensor<
        f32[24]
        EXLA.Backend
        [0.0421639047563076, 0.03801903501152992, 0.04582618921995163, 0.028266852721571922, 0.025787167251110077, 0.026369646191596985, 0.03499341756105423, 0.02089923992753029, 0.02381058968603611, 0.04512317106127739, 0.018083490431308746, 0.02251904457807541, 0.023619014769792557, 0.005619029048830271, 0.04942920058965683, 0.02654757723212242, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck8_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[384][1][3][3]
        EXLA.Backend
        [
          [
            [
              [-0.0020327651873230934, -0.03829871118068695, -0.013563335873186588],
              [-0.07023981958627701, -0.0530773401260376, -0.05999664217233658],
              [-0.02172546088695526, -0.05160605534911156, -0.02340804785490036]
            ]
          ],
          [
            [
              [-0.06414368003606796, -0.09891236573457718, 0.09147416055202484],
              [-0.035863131284713745, -0.026917265728116035, 0.07736452668905258],
              [0.08244364708662033, 0.10789858549833298, -0.16784991323947906]
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck11_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[576][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.0561985969543457, -0.08534234762191772, -0.0728735700249672],
              [0.20427849888801575, -0.03384411707520485, -0.18976956605911255],
              [0.10962294042110443, 0.006988390814512968, -0.03254517540335655]
            ]
          ],
          [
            [
              [-0.007281317375600338, -0.10424316674470901, -0.0011419359361752868],
              [-0.0659799799323082, -0.2147035449743271, -0.10787361860275269],
              [0.004407626111060381, 0.44921964406967163, ...]
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck14_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[960][160][1][1]
        EXLA.Backend
        [
          [
            [
              [0.037727002054452896]
            ],
            [
              [-0.07920914888381958]
            ],
            [
              [-0.0031863711774349213]
            ],
            [
              [-0.017933860421180725]
            ],
            [
              [-0.012083934620022774]
            ],
            [
              [-0.03305855393409729]
            ],
            [
              [0.010104630142450333]
            ],
            [
              [-0.01614421419799328]
            ],
            [
              [0.0415346585214138]
            ],
            [
              [-0.05910014733672142]
            ],
            [
              [3.380556881893426e-4]
            ],
            [
              [0.02735193446278572]
            ],
            [
              [-0.0788668617606163]
            ],
            [
              [-0.028101010248064995]
            ],
            [
              [-0.01531447283923626]
            ],
            [
              [-0.04052167013287544]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck15_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[960][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.008491688407957554, 0.01882096938788891, -0.001548346015624702],
              [0.009763970039784908, 0.2359771430492401, 0.03868749365210533],
              [-0.05791253224015236, 0.028106149286031723, -0.07315897196531296]
            ]
          ],
          [
            [
              [-0.01212818082422018, -0.16292721033096313, -0.033156707882881165],
              [0.005929793696850538, 0.2561517059803009, 0.014003193937242031],
              ...
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck7_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[64][384][1][1]
        EXLA.Backend
        [
          [
            [
              [-0.03408721089363098]
            ],
            [
              [-0.017362290993332863]
            ],
            [
              [0.01995295286178589]
            ],
            [
              [-0.051705457270145416]
            ],
            [
              [0.0689198449254036]
            ],
            [
              [0.08685684949159622]
            ],
            [
              [-0.09933704882860184]
            ],
            [
              [0.07178854197263718]
            ],
            [
              [0.044526729732751846]
            ],
            [
              [-0.018590005114674568]
            ],
            [
              [0.013196618296205997]
            ],
            [
              [-0.0795169249176979]
            ],
            [
              [-0.10800802707672119]
            ],
            [
              [-0.042785294353961945]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck4_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[192][1][3][3]
        EXLA.Backend
        [
          [
            [
              [-0.031621839851140976, 0.18290314078330994, 0.040609210729599],
              [-0.19228996336460114, 0.031516727060079575, 0.09460115432739258],
              [-0.026457147672772408, 0.05563095211982727, 0.0026122310664504766]
            ]
          ],
          [
            [
              [-0.00655049504712224, -0.004564541857689619, -0.01380461547523737],
              [-0.002132135210558772, ...],
              ...
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck16_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[960][160][1][1]
        EXLA.Backend
        [
          [
            [
              [0.04804954305291176]
            ],
            [
              [-0.0361880324780941]
            ],
            [
              [0.09869731217622757]
            ],
            [
              [0.03218498453497887]
            ],
            [
              [-0.05613848939538002]
            ],
            [
              [0.03235843405127525]
            ],
            [
              [0.035072293132543564]
            ],
            [
              [0.04535597562789917]
            ],
            [
              [0.031406670808792114]
            ],
            [
              [0.003971233498305082]
            ],
            [
              [0.016072509810328484]
            ],
            [
              [-0.041956160217523575]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck10_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [2.1335902999908285e-7, 3.535939754328865e-7, 5.073925990473072e-7, 5.08670098042785e-7, 3.3390904263796983e-7, 4.946827516505437e-7, 4.795706445293035e-7, 3.743390948329761e-7, 5.272184466775798e-7, 5.93354513966915e-7, 2.1419212714590685e-7, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [0.2979993522167206, 0.3140060007572174, 0.2445983588695526, 0.2865447998046875, 0.35691216588020325, 0.3592449426651001, 0.27936986088752747, 0.25922414660453796, 0.2633695900440216, 0.2901868522167206, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [-0.029136953875422478, -0.5148138999938965, 0.05061722546815872, -0.3127302825450897, -0.029322732239961624, 0.2176133394241333, 0.10727949440479279, -0.23169416189193726, -0.4668867290019989, ...]
      >,
      "var" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [0.11043474823236465, 0.09818199276924133, 0.06050458922982216, 0.09404341876506805, 0.19909954071044922, 0.13264349102973938, 0.06298568099737167, 0.07344550639390945, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck14_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[960][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.03072543255984783, 0.024941517040133476, -0.044987138360738754],
              [0.2105334848165512, -0.09797866642475128, -0.12514935433864594],
              [0.049303051084280014, -0.004124515224248171, -0.0421803742647171]
            ]
          ],
          [
            [
              [0.027953263372182846, ...],
              ...
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck12_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[576][96][1][1]
        EXLA.Backend
        [
          [
            [
              [-0.020311787724494934]
            ],
            [
              [0.00342985475435853]
            ],
            [
              [-0.04346230626106262]
            ],
            [
              [-0.016626032069325447]
            ],
            [
              [-0.012114123441278934]
            ],
            [
              [-0.012785198166966438]
            ],
            [
              [0.056096695363521576]
            ],
            [
              [0.09021388739347458]
            ],
            [
              [0.16139942407608032]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck0_batchnorm1_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[32]
        EXLA.Backend
        [1.037217140197754, 1.2865849733352661, 0.7645959258079529, 0.6374938488006592, 1.069508671760559, 0.6837575435638428, 0.9550023674964905, 0.9839419722557068, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[32]
        EXLA.Backend
        [0.2217620462179184, 0.11614178121089935, 0.1586529165506363, 0.15199987590312958, 0.23966756463050842, 0.17114117741584778, 0.18509696424007416, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[32]
        EXLA.Backend
        [0.6259635090827942, -0.08774495124816895, 1.0720069408416748, 2.3868372440338135, -5.589635848999023, 3.4053637981414795, ...]
      >,
      "var" => #Nx.Tensor<
        f32[32]
        EXLA.Backend
        [0.3122387230396271, 0.008777356706559658, 0.010774576105177402, 0.3379133641719818, 2.11269211769104, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck6_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[192][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.014363789930939674, 0.13237811625003815, 0.09334706515073776],
              [0.12919868528842926, -0.35225632786750793, -0.1466313749551773],
              [0.06540388613939285, ...]
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck16_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[320]
        EXLA.Backend
        [-5.189255602999765e-7, -2.7314237627251714e-7, -2.1782213366350334e-7, 1.9330661871208576e-7, 2.6874428016299134e-8, -2.4903325623881756e-8, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[320]
        EXLA.Backend
        [0.25196632742881775, 0.2486555427312851, 0.24830971658229828, 0.24681176245212555, 0.24941202998161316, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[320]
        EXLA.Backend
        [-0.1392788290977478, -0.05395103245973587, -0.14208653569221497, 0.18811774253845215, ...]
      >,
      "var" => #Nx.Tensor<
        f32[320]
        EXLA.Backend
        [0.008379261940717697, 0.008914480917155743, 0.007564138621091843, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck10_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[96][384][1][1]
        EXLA.Backend
        [
          [
            [
              [0.01824713684618473]
            ],
            [
              [0.08710300177335739]
            ],
            [
              [-0.10191874951124191]
            ],
            [
              [-0.07710960507392883]
            ],
            [
              [-0.042343202978372574]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck8_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[384][64][1][1]
        EXLA.Backend
        [
          [
            [
              [0.015875784680247307]
            ],
            [
              [0.00680553587153554]
            ],
            [
              [0.09295158088207245]
            ],
            [
              [-0.05411294475197792]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[32][3][3][3]
        EXLA.Backend
        [
          [
            [
              [-0.11137320101261139, -0.22142910957336426, 0.10717468708753586],
              ...
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck11_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [3.9526614159512974e-7, 5.761955890193349e-7, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [0.21182486414909363, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [...]
      >,
      ...
    },
    "mobilenetv20_features_linearbottleneck6_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[64][192][1][1]
        EXLA.Backend
        [
          [
            [
              [0.07365027070045471]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck3_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[32][144][1][1]
        EXLA.Backend
        [
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck4_conv0_fwd" => %{...},
    ...
  }
}
model
|> Axon.sigmoid()
|> Axon.Loop.evaluator()
|> Axon.Loop.metric(:accuracy)
|> Axon.Loop.run(test_pipeline, trained_model_state, compiler: EXLA)

08:16:14.945 [debug] Forwarding options: [compiler: EXLA] to JIT compiler
Batch: 22, accuracy: 0.9592391
%{
  0 => %{
    "accuracy" => #Nx.Tensor<
      f32
      EXLA.Backend
      0.95923912525177
    >
  }
}
model = model |> Axon.unfreeze(up: 50)

loss =
  &amp;Axon.Losses.binary_cross_entropy(&amp;1, &amp;2,
    reduction: :mean,
    from_logits: true
  )

optimizer = Polaris.Optimizers.rmsprop(learning_rate: 1.0e-5)

trained_model_state =
  model
  |> Axon.Loop.trainer(loss, optimizer)
  |> Axon.Loop.metric(:accuracy)
  |> Axon.Loop.validate(model, val_pipeline)
  |> Axon.Loop.early_stop("validation_loss", mode: :min, patience: 5)
  |> Axon.Loop.run(
    train_pipeline,
    trained_model_state,
    epochs: 100,
    compiler: EXLA
  )

08:17:43.369 [debug] Forwarding options: [compiler: EXLA] to JIT compiler
Epoch: 0, Batch: 700, accuracy: 0.9254190 loss: 0.1792496
Epoch: 1, Batch: 700, accuracy: 0.9409326 loss: 0.1616345
Epoch: 2, Batch: 700, accuracy: 0.9496701 loss: 0.1484112
Epoch: 3, Batch: 700, accuracy: 0.9557329 loss: 0.1378420
Epoch: 4, Batch: 700, accuracy: 0.9598787 loss: 0.1302559
Epoch: 5, Batch: 700, accuracy: 0.9650499 loss: 0.1230860
Epoch: 6, Batch: 700, accuracy: 0.9674572 loss: 0.1170269
Epoch: 7, Batch: 700, accuracy: 0.9712464 loss: 0.1114469
Batch: 6, accuracy: 0.9642857 loss: 0.1175858
...
        EXLA.Backend
        [0.1144828125834465, 0.29103899002075195, 0.19599959254264832, 0.11088918149471283, 0.11658263951539993, 0.2404186874628067, 0.05210668221116066, 0.279130220413208, 0.16053666174411774, 0.10330193489789963, 0.17037393152713776, 0.42853739857673645, 0.2756047546863556, 0.3480927646160126, 0.41298505663871765, 0.3465864956378937, 0.2918163239955902, 0.2078855335712433, 0.19570542871952057, 0.3232663571834564, 0.20460295677185059, 0.21620818972587585, 0.06636220216751099, 0.16237126290798187, 0.11230707168579102, 0.20659472048282623, 0.04134644940495491, 0.16707743704319, 0.18970169126987457, 0.08170202374458313, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [7.743461173959076e-4, -3.06177040329203e-4, -9.293855691794306e-5, -7.002947677392513e-5, 2.6267904104315676e-5, -3.16032295813784e-4, 1.1480091779958457e-4, 1.50591236888431e-4, 4.538445209618658e-4, -4.4312034151516855e-4, -1.6169094305951148e-4, 1.3050729467067868e-4, 7.128199649741873e-5, 5.354180466383696e-4, 0.0017610888462513685, -3.993306600023061e-4, -1.130728742282372e-5, 7.246735040098429e-4, 4.059264902025461e-4, -2.736354508670047e-5, 1.5211454592645168e-4, 1.1162958981003612e-4, -1.6184523701667786e-4, 4.564085975289345e-4, 1.0321872105123475e-4, -6.690184090984985e-5, 1.566050632391125e-4, 3.648195706773549e-5, -4.4106293353252113e-4, ...]
      >,
      "var" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [0.09844373911619186, 0.1737949252128601, 0.3841856122016907, 0.17619512975215912, 0.19387663900852203, 0.10076341778039932, 0.18362277746200562, 0.20341406762599945, 0.09384425729513168, 0.0681135356426239, 0.14232666790485382, 0.4347239136695862, 0.1753973662853241, 0.4196370542049408, 0.3148811459541321, 0.3903367221355438, 0.0904870554804802, 0.1093224361538887, 0.33561405539512634, 0.118010513484478, 0.1954060047864914, 0.1988266259431839, 0.13409970700740814, 0.26589640974998474, 0.16766609251499176, 0.1590973138809204, 0.09238500148057938, 0.27864906191825867, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck11_batchnorm0_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [0.1379205733537674, 0.1852763593196869, 0.07007614523172379, 0.14786551892757416, 0.08499550819396973, 0.022980524227023125, 0.10445888340473175, -0.03380206599831581, 0.05621454492211342, 0.13879631459712982, 0.1278844028711319, 0.11503273248672485, 0.008648642338812351, 0.07986151427030563, 0.05827014893293381, 0.02635921724140644, 0.16676078736782074, 0.1696624606847763, -0.0049711973406374454, 0.11680029332637787, 0.09153503179550171, -0.09545790404081345, 0.10052905976772308, 0.2198638916015625, 0.09066508710384369, 0.021357741206884384, 0.1243291050195694, 0.13527752459049225, 0.09010029584169388, 0.07285095006227493, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [0.1116461232304573, 0.08017633110284805, 0.11955763399600983, 0.09390335530042648, 0.09855279326438904, 0.14738331735134125, 0.0888434574007988, 0.13319958746433258, 0.10988157987594604, 0.09758485853672028, 0.10826665163040161, 0.1016610860824585, 0.09888888150453568, 0.12156049907207489, 0.11694491654634476, 0.119488425552845, 0.07615012675523758, 0.06269187480211258, 0.15567438304424286, 0.1064029410481453, 0.11337470263242722, 0.14443783462047577, 0.0840112492442131, 0.09276950359344482, 0.16790293157100677, 0.11844969540834427, 0.0930234044790268, 0.11310834437608719, 0.1162751242518425, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [3.63793503765919e-7, 1.209602089602413e-7, 5.4800288751266635e-8, 2.552324076532386e-7, 3.2612440747925575e-8, 1.4689121030642127e-8, 1.1114344999896275e-7, -8.35849434110969e-8, 2.2601201976613083e-7, 1.6091770760340296e-7, 3.1186812066152925e-7, -1.1275391287313141e-8, -2.748528871165945e-8, -1.7028881416081276e-7, -1.3741856719207135e-7, -1.3034306789450056e-7, 4.320734205975896e-7, 5.223997732173302e-7, 1.3161741918565895e-8, 1.749855726984606e-7, 6.074458980265263e-8, 4.904541128780693e-7, 2.8845341670802327e-9, 3.752578692228781e-8, -8.265797646345163e-8, -2.6385482243540537e-8, 6.387549156272598e-9, -3.501408230022207e-8, ...]
      >,
      "var" => #Nx.Tensor<
        f32[576]
        EXLA.Backend
        [0.02586364559829235, 0.04419300705194473, 0.02695547789335251, 0.027057670056819916, 0.023040946573019028, 0.03236440196633339, 0.021965283900499344, 0.03861037641763687, 0.02072414942085743, 0.020027386024594307, 0.017105598002672195, 0.027752427384257317, 0.025669051334261894, 0.030634647235274315, 0.02309420332312584, 0.026124022901058197, 0.009496251121163368, 0.02531951665878296, 0.04465075209736824, 0.03168340027332306, 0.03208788484334946, 0.03684369474649429, 0.01707514189183712, 0.033794183284044266, 0.05672197788953781, 0.01841931976377964, 0.031393010169267654, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck15_batchnorm1_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [-0.0816231295466423, 0.12648126482963562, -0.007565976586192846, -0.04337223619222641, 0.07973095774650574, 0.00586037989705801, -0.04332118108868599, -0.062439367175102234, 0.07367430627346039, 0.008881223388016224, -0.020015940070152283, 0.03910618647933006, -0.05755405128002167, -0.04922578111290932, -0.07715956121683121, 0.07127390801906586, 0.009607778862118721, -0.0018916352419182658, -0.06841062009334564, 0.03574378415942192, -0.06820035725831985, -0.08501345664262772, -0.1458040475845337, -0.0010745554463937879, -0.014342773705720901, -0.07032705098390579, 0.05578082799911499, 0.34916213154792786, -0.032628338783979416, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [0.31593403220176697, 0.16369463503360748, 0.2266220897436142, 0.23600031435489655, 0.16440345346927643, 0.14116989076137543, 0.19827400147914886, 0.15404096245765686, 0.14843463897705078, 0.13872063159942627, 0.18212608993053436, 0.19958379864692688, 0.20408156514167786, 0.2203507125377655, 0.22687706351280212, 0.15941059589385986, 0.13872721791267395, 0.13495267927646637, 0.200734481215477, 0.17756947875022888, 0.1950230598449707, 0.18064196407794952, 0.2480158656835556, 0.14279963076114655, 0.29032111167907715, 0.21508155763149261, 0.16373638808727264, 0.11638444662094116, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [0.006607928778976202, -0.005489333998411894, 0.001202044659294188, 0.016741983592510223, -0.021014627069234848, 3.9149439544416964e-4, -0.001380614354275167, 0.006143252365291119, -0.022377200424671173, 3.921858442481607e-4, 0.011808911338448524, -0.038100697100162506, -0.0027107521891593933, -0.002684985054656863, -0.003361391369253397, -0.027487538754940033, 0.0032275188714265823, 7.377289002761245e-4, -0.003101653652265668, -0.02393931709229946, -0.0026154580991715193, 4.0451946551911533e-4, 0.019924011081457138, 0.0014483053237199783, 0.012146278284490108, -6.006719195283949e-4, -0.02536875568330288, ...]
      >,
      "var" => #Nx.Tensor<
        f32[960]
        EXLA.Backend
        [6.60044839605689e-4, 2.093793300446123e-4, 6.339335814118385e-4, 8.293294231407344e-4, 4.367101937532425e-4, 3.0174001949490048e-5, 1.974338520085439e-4, 5.38308268005494e-5, 5.038247327320278e-4, 1.9694401999004185e-4, 3.8651813520118594e-4, 5.852532922290266e-4, 2.917999809142202e-4, 3.8506241980940104e-4, 3.85848106816411e-4, 4.808907106053084e-4, 8.100475679384544e-5, 6.1370710682240315e-6, 1.8867437029257417e-4, 3.83470905944705e-4, 5.22851652931422e-4, 4.0124301449395716e-5, 4.517717461567372e-4, 7.429129254887812e-6, 0.0013210318284109235, 2.3677274293731898e-4, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck9_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[64][384][1][1]
        EXLA.Backend
        [
          [
            [
              [-0.029621589928865433]
            ],
            [
              [0.03196175768971443]
            ],
            [
              [-0.04775437340140343]
            ],
            [
              [-0.02047593705356121]
            ],
            [
              [0.022681893780827522]
            ],
            [
              [0.0061656758189201355]
            ],
            [
              [-0.020124075934290886]
            ],
            [
              [-0.0419141948223114]
            ],
            [
              [-0.0958712100982666]
            ],
            [
              [-0.03112773969769478]
            ],
            [
              [0.033005211502313614]
            ],
            [
              [-0.0796334370970726]
            ],
            [
              [-0.08826649934053421]
            ],
            [
              [0.008132039569318295]
            ],
            [
              [0.13005898892879486]
            ],
            [
              [0.01112390961498022]
            ],
            [
              [-0.025056127458810806]
            ],
            [
              [-0.021656425669789314]
            ],
            [
              [-0.009619093500077724]
            ],
            [
              [-0.09488093852996826]
            ],
            [
              [-0.03190577030181885]
            ],
            [
              [0.08531086146831512]
            ],
            [
              [0.047572940587997437]
            ],
            [
              [-0.036207523196935654]
            ],
            [
              [-0.1315651535987854]
            ],
            [
              [-0.03746785223484039]
            ],
            [
              [-0.04043968394398689]
            ],
            [
              [-0.1008371040225029]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck11_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[96][576][1][1]
        EXLA.Backend
        [
          [
            [
              [-0.04216459393501282]
            ],
            [
              [-0.03359765186905861]
            ],
            [
              [0.021302534267306328]
            ],
            [
              [0.06966963410377502]
            ],
            [
              [-0.023086098954081535]
            ],
            [
              [0.02423691935837269]
            ],
            [
              [-0.0027191585395485163]
            ],
            [
              [-0.12812955677509308]
            ],
            [
              [0.10142464935779572]
            ],
            [
              [0.04367773234844208]
            ],
            [
              [-0.04808792844414711]
            ],
            [
              [-0.059124793857336044]
            ],
            [
              [-0.03544275462627411]
            ],
            [
              [-0.10292188823223114]
            ],
            [
              [-0.15540222823619843]
            ],
            [
              [0.08993023633956909]
            ],
            [
              [-0.0018712839810177684]
            ],
            [
              [0.04499407112598419]
            ],
            [
              [0.03717062994837761]
            ],
            [
              [0.1088862419128418]
            ],
            [
              [0.03789613023400307]
            ],
            [
              [0.011467562057077885]
            ],
            [
              [0.022153226658701897]
            ],
            [
              [0.011087901890277863]
            ],
            [
              [0.10830669850111008]
            ],
            [
              [-0.039926622062921524]
            ],
            [
              [0.03126271069049835]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck2_batchnorm1_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[144]
        EXLA.Backend
        [0.06266669183969498, 0.34174367785453796, 0.014498202130198479, 0.09171395748853683, -0.01248596515506506, 0.09248802065849304, 0.03913332521915436, -0.12479757517576218, 0.04716098681092262, 0.2623796761035919, 0.00875355675816536, 0.3866884410381317, 0.34793952107429504, 0.07759835571050644, 0.28590700030326843, 0.34494516253471375, 0.1602264642715454, 0.3233429491519928, -0.08500878512859344, 0.017791876569390297, 0.012547487393021584, 0.07443935424089432, 0.470037043094635, 0.06998748332262039, -0.1145705133676529, 0.35432684421539307, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[144]
        EXLA.Backend
        [0.14569389820098877, 0.08483350276947021, 0.18709848821163177, 0.08500487357378006, 0.0695008710026741, 0.13266052305698395, 0.23066304624080658, 0.1294335573911667, 0.15951289236545563, 0.1612233966588974, 0.04986649379134178, 0.24220621585845947, 0.13059896230697632, 0.2140040248632431, 0.16401590406894684, 0.04375945031642914, 0.1284344643354416, 0.10862768441438675, 0.22237107157707214, 0.20400500297546387, 0.2341395914554596, 0.36944401264190674, 0.1626269817352295, 0.1425088793039322, 0.37626516819000244, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[144]
        EXLA.Backend
        [0.04842672124505043, 0.011466434225440025, -0.28272300958633423, 0.04086904227733612, -0.0, 0.03728378936648369, -0.07810080051422119, -0.0095414062961936, 0.019914837554097176, -0.26291584968566895, 1.200515914815965e-22, 0.4378272294998169, -0.07176417112350464, -0.20427647233009338, -0.40337827801704407, 0.07497201859951019, 0.06511131674051285, 0.02972135879099369, -0.029825769364833832, 0.03816595301032066, 0.008758057840168476, -0.5172790288925171, 0.18057824671268463, 0.03219117224216461, ...]
      >,
      "var" => #Nx.Tensor<
        f32[144]
        EXLA.Backend
        [0.0016082772053778172, 0.008874273858964443, 0.014802338555455208, 0.010124930180609226, 0.0, 0.0027794779743999243, 0.019753890112042427, 6.709994049742818e-4, 0.004814377520233393, 0.018505984917283058, 3.875728477835647e-26, 0.024166392162442207, 0.011687243357300758, 0.018400687724351883, 0.00805666297674179, 0.003925488330423832, 0.004267431795597076, 0.005561113357543945, 0.00586837250739336, 0.01074590440839529, 0.023998573422431946, 0.010233975946903229, 0.007847330532968044, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck3_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[144][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.007945280522108078, -0.5135775208473206, 0.642527163028717],
              [0.09571424126625061, -0.4191504120826721, -0.42046400904655457],
              [0.2738749086856842, -1.341597080230713, 1.7941224575042725]
            ]
          ],
          [
            [
              [0.36698344349861145, 0.8006415963172913, 1.0355921983718872],
              [0.2570298910140991, 0.9058852791786194, 0.1435706615447998],
              [0.08509054780006409, -0.29565608501434326, -0.19910597801208496]
            ]
          ],
          [
            [
              [-0.07826036959886551, 1.5095280408859253, -0.7258926033973694],
              [-0.1067059189081192, -0.07988816499710083, -0.4704415500164032],
              [-2.0861942768096924, ...]
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck10_batchnorm0_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [-0.14200253784656525, -0.08138399571180344, 0.020918697118759155, -0.030422303825616837, -0.054856158792972565, 0.0021326385904103518, -0.18490885198116302, -0.03978404402732849, 0.17561353743076324, 0.007802810054272413, -0.05954292416572571, 0.001685986528173089, -0.021205948665738106, 0.09003060311079025, 0.03726910054683685, 0.14319199323654175, 0.15906931459903717, 0.09075694531202316, 0.14122611284255981, 0.19317775964736938, 0.11409415304660797, 0.2862238883972168, 0.21863330900669098, 0.31881415843963623, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [0.08955207467079163, 0.12251311540603638, 0.13474930822849274, 0.12169622629880905, 0.10408146679401398, 0.12673817574977875, 0.09429759532213211, 0.004978157579898834, 0.06450905650854111, 0.12292398512363434, 0.008911419659852982, 0.16245806217193604, 0.12929747998714447, 0.07689117640256882, 0.13534259796142578, 0.0939403846859932, 0.07975732535123825, 0.07448924332857132, 0.02857014909386635, 0.05634821206331253, 0.04961058497428894, 0.08511262387037277, 0.1229783222079277, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [2.306738606421277e-4, 8.843086543492973e-4, 2.0830532594118267e-4, 2.6212024386040866e-4, -1.6121921362355351e-4, 3.321829717606306e-4, 1.4777167234569788e-4, 8.781474753050134e-5, -5.305309314280748e-4, 3.368522447999567e-4, 1.637199311517179e-4, 8.147450280375779e-4, 4.6918101725168526e-4, 2.0563749421853572e-4, 7.5969856698066e-4, -1.2368688476271927e-4, -2.2818840807303786e-4, 1.4898755762260407e-4, -1.767556750564836e-5, 1.1803162487922236e-4, 3.7923167110420763e-4, 8.40145457914332e-6, ...]
      >,
      "var" => #Nx.Tensor<
        f32[384]
        EXLA.Backend
        [0.03833653777837753, 0.19651877880096436, 0.16792143881320953, 0.121506467461586, 0.07553654164075851, 0.15889404714107513, 0.06837835162878036, 0.0039282916113734245, 0.33895161747932434, 0.1504734307527542, 0.011571435257792473, 0.34775444865226746, 0.1949744075536728, 0.12849894165992737, 0.28149864077568054, 0.12852248549461365, 0.1332881599664688, 0.09948007017374039, 0.1221485584974289, 0.20862258970737457, 0.12218637019395828, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck14_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[160]
        EXLA.Backend
        [-1.6225178001150198e-7, -1.8454430517067522e-7, -8.0369410682124e-8, 4.6417799381970326e-8, -7.48215001067365e-8, -2.8193767320772167e-7, 7.217398945158493e-8, -8.68730083425362e-8, 9.796455913146929e-8, 3.460570496827131e-8, -9.80068719513838e-8, -5.421593129995017e-8, 2.481537819676305e-7, 1.395186473018839e-7, 4.2195118510335305e-9, -2.3262117565536755e-7, 2.2320443804346723e-7, 3.619721056224989e-8, -1.2108348812489567e-7, -9.21366449802008e-9, 5.413741988036236e-8, 1.0924743065743314e-7, -2.9503431164812355e-7, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[160]
        EXLA.Backend
        [0.1874222606420517, 0.3471328020095825, 0.24452398717403412, 0.33153024315834045, 0.21980319917201996, 0.3147669732570648, 0.2848625183105469, 0.16946527361869812, 0.3173196613788605, 0.227366104722023, 0.21156655251979828, 0.22863364219665527, 0.22919036448001862, 0.17270319163799286, 0.27297401428222656, 0.22551169991493225, 0.2640722990036011, 0.1800961047410965, 0.25774723291397095, 0.22528624534606934, 0.24925146996974945, 0.22121092677116394, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[160]
        EXLA.Backend
        [-0.10289469361305237, 0.13421736657619476, 0.3127436637878418, -0.010098671540617943, 0.11501706391572952, 0.07170626521110535, -0.12338873744010925, -0.09662780910730362, 0.10249986499547958, 0.039654918015003204, 0.23144079744815826, -0.02762533910572529, -0.2323276251554489, -0.14802084863185883, -0.17405559122562408, -0.05928152799606323, 0.30801624059677124, 0.03448152542114258, 0.10294467210769653, -0.14233750104904175, 0.042918454855680466, ...]
      >,
      "var" => #Nx.Tensor<
        f32[160]
        EXLA.Backend
        [0.01982937380671501, 0.07223469018936157, 0.03559407591819763, 0.05254514887928963, 0.03399147093296051, 0.05125325173139572, 0.045940618962049484, 0.015112430788576603, 0.0665295273065567, 0.029054192826151848, 0.024110015481710434, 0.028089039027690887, 0.032291218638420105, 0.017801793292164803, 0.04934801161289215, 0.028767650946974754, 0.036390747874975204, 0.020782455801963806, 0.038260187953710556, 0.03192899003624916, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck0_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[16]
        EXLA.Backend
        [-0.0010560675291344523, -6.787155289202929e-4, -7.046025712043047e-4, -9.398604743182659e-4, -8.966192253865302e-4, -2.11737904464826e-4, -5.149749922566116e-4, -0.0010324727045372128, -7.925425306893885e-4, -9.090156527236104e-5, -4.453391011338681e-4, -9.805577574297786e-4, -7.45584664400667e-4, -7.640189724043012e-4, -0.0011470100143924356, -3.8889769348315895e-4]
      >,
      "gamma" => #Nx.Tensor<
        f32[16]
        EXLA.Backend
        [0.5097491145133972, 0.6192775368690491, 0.657399594783783, 0.5558041930198669, 0.6164559721946716, 0.6381960511207581, 0.6152222156524658, 0.4946329891681671, 0.5616724491119385, 0.6567564606666565, 0.5526584982872009, 0.42007970809936523, 0.4637764096260071, 0.5626768469810486, 0.6189385056495667, 0.6354565620422363]
      >,
      "mean" => #Nx.Tensor<
        f32[16]
        EXLA.Backend
        [0.02103538066148758, -0.47714170813560486, -0.4624972641468048, 1.7340151071548462, 0.3934422731399536, -0.8457171320915222, 1.2445671558380127, -0.4817531406879425, 0.8393098711967468, -2.1945483684539795, 2.285907030105591, 0.04294494166970253, -1.1515527963638306, -0.3779143691062927, -3.0484728813171387, -0.784763514995575]
      >,
      "var" => #Nx.Tensor<
        f32[16]
        EXLA.Backend
        [0.015541340224444866, 0.057948190718889236, 0.0564122200012207, 0.01880471035838127, 0.03849662467837334, 0.09232877939939499, 0.053476687520742416, 0.04164041951298714, 0.03743227198719978, 0.14161214232444763, 0.04085685685276985, 0.0607682503759861, 0.034746140241622925, 0.05628001317381859, 0.04128635302186012, 0.0522528812289238]
      >
    },
    "mobilenetv20_features_linearbottleneck7_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[384][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.09069628268480301, 0.011556080542504787, -0.044405367225408554],
              [0.1974998563528061, -0.05144780874252319, -0.15406860411167145],
              [0.03294772282242775, -0.08055609464645386, -0.10401368886232376]
            ]
          ],
          [
            [
              [0.004530614707618952, -0.016135336831212044, 0.010837683454155922],
              [-0.013911036774516106, -0.17941752076148987, 0.0742872878909111],
              [0.03701423481106758, 0.17489252984523773, 0.049539677798748016]
            ]
          ],
          [
            [
              [-0.03442509472370148, -0.04338786378502846, -8.070958429016173e-4],
              ...
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck1_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[96][16][1][1]
        EXLA.Backend
        [
          [
            [
              [-0.26627230644226074]
            ],
            [
              [0.1223592683672905]
            ],
            [
              [-0.17675793170928955]
            ],
            [
              [-0.14175795018672943]
            ],
            [
              [-0.04847275838255882]
            ],
            [
              [0.02572016417980194]
            ],
            [
              [0.18247880041599274]
            ],
            [
              [-0.22272470593452454]
            ],
            [
              [-0.0375950001180172]
            ],
            [
              [-0.12036115676164627]
            ],
            [
              [-0.42806223034858704]
            ],
            [
              [0.2049194574356079]
            ],
            [
              [0.025416148826479912]
            ],
            [
              [-0.20868347585201263]
            ],
            [
              [-0.00869692862033844]
            ],
            [
              [0.10014521330595016]
            ]
          ],
          [
            [
              [0.1901395469903946]
            ],
            [
              [-0.15854698419570923]
            ],
            [
              [-0.05858785659074783]
            ],
            [
              [-0.016011832281947136]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck2_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[24]
        EXLA.Backend
        [-6.408389890566468e-5, -1.0128262510988861e-4, -4.3479518353706226e-5, -1.589354214956984e-4, -7.557999197160825e-5, -2.3196584152174182e-5, -2.216532493548584e-6, 9.525781933916733e-6, -8.600373985245824e-5, -2.3008385323919356e-4, -2.3174902889877558e-4, -2.012878976529464e-4, -9.838923870120198e-5, -1.22729703434743e-4, -9.646145917940885e-5, -2.797471643134486e-5, -1.039204653352499e-4, -1.2249314750079066e-4, -8.16335596027784e-5, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[24]
        EXLA.Backend
        [0.5065965056419373, 0.44786974787712097, 0.5534567832946777, 0.3781253695487976, 0.44674113392829895, 0.28593751788139343, 0.5029821395874023, 0.19650447368621826, 0.3918337821960449, 0.33050450682640076, 0.31709131598472595, 0.1402808576822281, 0.4198921024799347, 0.02611447498202324, 0.5337468981742859, 0.49816271662712097, 0.32386210560798645, 0.40961867570877075, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[24]
        EXLA.Backend
        [-0.49676281213760376, -0.5412184000015259, 0.2777172029018402, 0.30402812361717224, -0.038367077708244324, 0.19269412755966187, -0.20861922204494476, 0.3330950438976288, -0.35258007049560547, -0.0413212850689888, -0.44278526306152344, 0.43556883931159973, -0.007557356264442205, -0.049789782613515854, -0.1810888946056366, -0.1666330248117447, 0.15477608144283295, ...]
      >,
      "var" => #Nx.Tensor<
        f32[24]
        EXLA.Backend
        [0.0421639047563076, 0.03801903501152992, 0.04582618921995163, 0.028266852721571922, 0.025787167251110077, 0.026369646191596985, 0.03499341756105423, 0.02089923992753029, 0.02381058968603611, 0.04512317106127739, 0.018083490431308746, 0.02251904457807541, 0.023619014769792557, 0.005619029048830271, 0.04942920058965683, 0.02654757723212242, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck8_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[384][1][3][3]
        EXLA.Backend
        [
          [
            [
              [-0.0020327651873230934, -0.03829871118068695, -0.013563335873186588],
              [-0.07023981958627701, -0.0530773401260376, -0.05999664217233658],
              [-0.02172546088695526, -0.05160605534911156, -0.02340804785490036]
            ]
          ],
          [
            [
              [-0.06414368003606796, -0.09891236573457718, 0.09147416055202484],
              [-0.035863131284713745, -0.026917265728116035, 0.07736452668905258],
              [0.08244364708662033, 0.10789858549833298, -0.16784991323947906]
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck11_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[576][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.0561985969543457, -0.08534234762191772, -0.0728735700249672],
              [0.20427849888801575, -0.03384411707520485, -0.18976956605911255],
              [0.10962294042110443, 0.006988390814512968, -0.03254517540335655]
            ]
          ],
          [
            [
              [-0.007281317375600338, -0.10424316674470901, -0.0011419359361752868],
              [-0.0659799799323082, -0.2147035449743271, -0.10787361860275269],
              [0.004407626111060381, 0.44921964406967163, ...]
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck14_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[960][160][1][1]
        EXLA.Backend
        [
          [
            [
              [0.03726569190621376]
            ],
            [
              [-0.07747243344783783]
            ],
            [
              [-0.0027355963829904795]
            ],
            [
              [-0.018318168818950653]
            ],
            [
              [-0.011423363350331783]
            ],
            [
              [-0.033270787447690964]
            ],
            [
              [0.009682310745120049]
            ],
            [
              [-0.015885118395090103]
            ],
            [
              [0.04243874177336693]
            ],
            [
              [-0.05969354510307312]
            ],
            [
              [4.88394289277494e-4]
            ],
            [
              [0.028318190947175026]
            ],
            [
              [-0.07891203463077545]
            ],
            [
              [-0.029163353145122528]
            ],
            [
              [-0.016083136200904846]
            ],
            [
              [-0.04072346165776253]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck15_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[960][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.007893153466284275, 0.02030854858458042, -0.00212265201844275],
              [0.009393330663442612, 0.23728495836257935, 0.03770001605153084],
              [-0.05673065781593323, 0.028933633118867874, -0.07290589064359665]
            ]
          ],
          [
            [
              [-0.010928697884082794, -0.1602908819913864, -0.03171931579709053],
              [0.008520246483385563, 0.258930504322052, 0.015797391533851624],
              ...
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck7_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[64][384][1][1]
        EXLA.Backend
        [
          [
            [
              [-0.03408721089363098]
            ],
            [
              [-0.017362290993332863]
            ],
            [
              [0.01995295286178589]
            ],
            [
              [-0.051705457270145416]
            ],
            [
              [0.0689198449254036]
            ],
            [
              [0.08685684949159622]
            ],
            [
              [-0.09933704882860184]
            ],
            [
              [0.07178854197263718]
            ],
            [
              [0.044526729732751846]
            ],
            [
              [-0.018590005114674568]
            ],
            [
              [0.013196618296205997]
            ],
            [
              [-0.0795169249176979]
            ],
            [
              [-0.10800802707672119]
            ],
            [
              [-0.042785294353961945]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck4_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[192][1][3][3]
        EXLA.Backend
        [
          [
            [
              [-0.031621839851140976, 0.18290314078330994, 0.040609210729599],
              [-0.19228996336460114, 0.031516727060079575, 0.09460115432739258],
              [-0.026457147672772408, 0.05563095211982727, 0.0026122310664504766]
            ]
          ],
          [
            [
              [-0.00655049504712224, -0.004564541857689619, -0.01380461547523737],
              [-0.002132135210558772, ...],
              ...
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck16_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[960][160][1][1]
        EXLA.Backend
        [
          [
            [
              [0.04690628498792648]
            ],
            [
              [-0.03638744354248047]
            ],
            [
              [0.10043051093816757]
            ],
            [
              [0.030933288857340813]
            ],
            [
              [-0.05413831025362015]
            ],
            [
              [0.0305971410125494]
            ],
            [
              [0.03439365327358246]
            ],
            [
              [0.044807445257902145]
            ],
            [
              [0.029652398079633713]
            ],
            [
              [0.0018540085293352604]
            ],
            [
              [0.01651659607887268]
            ],
            [
              [-0.041166260838508606]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck10_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [2.1335902999908285e-7, 3.535939754328865e-7, 5.073925990473072e-7, 5.08670098042785e-7, 3.3390904263796983e-7, 4.946827516505437e-7, 4.795706445293035e-7, 3.743390948329761e-7, 5.272184466775798e-7, 5.93354513966915e-7, 2.1419212714590685e-7, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [0.2979993522167206, 0.3140060007572174, 0.2445983588695526, 0.2865447998046875, 0.35691216588020325, 0.3592449426651001, 0.27936986088752747, 0.25922414660453796, 0.2633695900440216, 0.2901868522167206, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [-0.029722288250923157, -0.514879047870636, 0.05067770555615425, -0.3127911388874054, -0.02917076274752617, 0.21818682551383972, 0.1075483188033104, -0.231681227684021, -0.46620145440101624, ...]
      >,
      "var" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [0.10969194769859314, 0.09763853251934052, 0.060802262276411057, 0.09436583518981934, 0.19624005258083344, 0.13399700820446014, 0.06356355547904968, 0.07258738577365875, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck14_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[960][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.034780435264110565, 0.029322482645511627, -0.04238593950867653],
              [0.21440976858139038, -0.09431995451450348, -0.12132745236158371],
              [0.050924308598041534, -0.0021802415139973164, -0.039994433522224426]
            ]
          ],
          [
            [
              [0.02731008641421795, ...],
              ...
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck12_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[576][96][1][1]
        EXLA.Backend
        [
          [
            [
              [-0.020311787724494934]
            ],
            [
              [0.00342985475435853]
            ],
            [
              [-0.04346230626106262]
            ],
            [
              [-0.016626032069325447]
            ],
            [
              [-0.012114123441278934]
            ],
            [
              [-0.012785198166966438]
            ],
            [
              [0.056096695363521576]
            ],
            [
              [0.09021388739347458]
            ],
            [
              [0.16139942407608032]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck0_batchnorm1_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[32]
        EXLA.Backend
        [1.037217140197754, 1.2865849733352661, 0.7645959258079529, 0.6374938488006592, 1.069508671760559, 0.6837575435638428, 0.9550023674964905, 0.9839419722557068, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[32]
        EXLA.Backend
        [0.2217620462179184, 0.11614178121089935, 0.1586529165506363, 0.15199987590312958, 0.23966756463050842, 0.17114117741584778, 0.18509696424007416, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[32]
        EXLA.Backend
        [0.6259415745735168, -0.08773765712976456, 1.0720137357711792, 2.3871850967407227, -5.589552402496338, 3.4055421352386475, ...]
      >,
      "var" => #Nx.Tensor<
        f32[32]
        EXLA.Backend
        [0.31177079677581787, 0.008774759247899055, 0.010728973895311356, 0.33693572878837585, 2.11619234085083, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck6_conv1_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[192][1][3][3]
        EXLA.Backend
        [
          [
            [
              [0.014363789930939674, 0.13237811625003815, 0.09334706515073776],
              [0.12919868528842926, -0.35225632786750793, -0.1466313749551773],
              [0.06540388613939285, ...]
            ]
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck16_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[320]
        EXLA.Backend
        [-3.0071265655351453e-7, -3.118751124020491e-7, 3.883552324168704e-8, 3.05521155041788e-8, -1.3791427022624703e-7, -1.4933878844658466e-7, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[320]
        EXLA.Backend
        [0.25120291113853455, 0.24765633046627045, 0.2482394427061081, 0.2461756020784378, 0.2500693202018738, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[320]
        EXLA.Backend
        [-0.12367760390043259, -0.05487632751464844, -0.12968169152736664, 0.1629655510187149, ...]
      >,
      "var" => #Nx.Tensor<
        f32[320]
        EXLA.Backend
        [0.009201690554618835, 0.011435198597609997, 0.00852434802800417, ...]
      >
    },
    "mobilenetv20_features_linearbottleneck10_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[96][384][1][1]
        EXLA.Backend
        [
          [
            [
              [0.01824713684618473]
            ],
            [
              [0.08710300177335739]
            ],
            [
              [-0.10191874951124191]
            ],
            [
              [-0.07710960507392883]
            ],
            [
              [-0.042343202978372574]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck8_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[384][64][1][1]
        EXLA.Backend
        [
          [
            [
              [0.015875784680247307]
            ],
            [
              [0.00680553587153554]
            ],
            [
              [0.09295158088207245]
            ],
            [
              [-0.05411294475197792]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_conv0_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[32][3][3][3]
        EXLA.Backend
        [
          [
            [
              [-0.11137320101261139, -0.22142910957336426, 0.10717468708753586],
              ...
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck11_batchnorm2_fwd" => %{
      "beta" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [3.9526614159512974e-7, 5.761955890193349e-7, ...]
      >,
      "gamma" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [0.21182486414909363, ...]
      >,
      "mean" => #Nx.Tensor<
        f32[96]
        EXLA.Backend
        [...]
      >,
      ...
    },
    "mobilenetv20_features_linearbottleneck6_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[64][192][1][1]
        EXLA.Backend
        [
          [
            [
              [0.07365027070045471]
            ],
            ...
          ],
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck3_conv2_fwd" => %{
      "kernel" => #Nx.Tensor<
        f32[32][144][1][1]
        EXLA.Backend
        [
          ...
        ]
      >
    },
    "mobilenetv20_features_linearbottleneck4_conv0_fwd" => %{...},
    ...
  }
}
model
|> Axon.sigmoid()
|> Axon.Loop.evaluator()
|> Axon.Loop.metric(:accuracy)
|> Axon.Loop.run(test_pipeline, trained_model_state, compiler: EXLA)

16:58:36.174 [debug] Forwarding options: [compiler: EXLA] to JIT compiler
Batch: 22, accuracy: 0.9755435
%{
  0 => %{
    "accuracy" => #Nx.Tensor<
      f32
      EXLA.Backend
      0.9755434989929199
    >
  }
}