public-apps/lab1.livemd

<!-- livebook:{"persist_outputs":true} -->

# Cleaned vs Dirty V2

```elixir
Mix.install(
  [
    {:stb_image, "~> 0.5.2"},
    {:axon, "~> 0.5"},
    {:polaris, "~> 0.1"},
    {:exla, "~> 0.5"},
    {:explorer, "~> 0.8.0"},
    {:nx, "~> 0.5"},
    {:kino_ripmd, github: "clm-a/kino_ripmd"},
    {:kino_explorer, "~> 0.1.0"},
    {:kino_vega_lite, "~> 0.1.0"},
    {:kino, "~> 0.10.0"},
    {:csv, "~> 3.2"}
  ],
  config: [
    nx: [
      default_backend: EXLA.Backend,
      default_defn_options: [compiler: EXLA]
    ],
    exla: [
      default_client: :cuda,
      clients: [
        cuda: [platform: :cuda],
        rocm: [platform: :rocm],
        tpu: [platform: :tpu],
        host: [platform: :host]
      ],
      memory_fraction: 0.9,
      preallocate: false
    ]
  ],
  system_env: [
    XLA_TARGET: "cuda120"
  ]
)
```

## Goal

_Hi! It is boring to wash the dishes. Luckily, half of them are already clean. Train a classifier to determine clean ones to save time for the new machine learning course ;)_

_It is a few shot learning competition. We have a dataset of 20 clean and 20 dirty plates in train and hundreds of plates in test. Good luck!_

Igor.Slinko. (2019). Cleaned vs Dirty V2. Kaggle. https://kaggle.com/competitions/platesv2

## Setting up the model

```elixir
alias Axon.Loop.State
import Nx.Defn

directories =
  "/home/le-moski/Documents/FEFU/7S/AI/lab1-dirtyplates/platesv2/plates/train/{cleaned,dirty}/*.jpg"

batch_size = 8
image_channels = 3
image_w = 256
image_h = 256
channel_value_max = 255

cleaned_class = Nx.tensor([1, 0], type: {:u, 8})
dirty_class = Nx.tensor([0, 1], type: {:u, 8})
```

<!-- livebook:{"output":true} -->

```
#Nx.Tensor<
  u8[2]
  EXLA.Backend<cuda:0, 0.463371035.4063887383.206028>
  [0, 1]
>
```

```elixir
parse_img = fn filename ->
  class =
    if Path.dirname(filename)
       |> String.split("/")
       |> List.last() == "cleaned" do
      cleaned_class
    else
      dirty_class
    end

  {:ok, img} = StbImage.read_file(filename)

  img = StbImage.resize(img, image_h, image_w)

  {StbImage.to_nx(img), class}
end
```

<!-- livebook:{"output":true} -->

```
#Function<42.39164016/1 in :erl_eval.expr/6>
```

```elixir
data =
  Path.wildcard(directories)
  |> Enum.shuffle()
  |> Stream.chunk_every(batch_size, batch_size)
  |> Task.async_stream(fn batch ->
    {imgs, classes} = batch |> Enum.map(&parse_img.(&1)) |> Enum.unzip()
    {Nx.stack(imgs), Nx.stack(classes)}
  end)
  |> Stream.map(fn {:ok, {imgs, classes}} -> {imgs |> Nx.divide(channel_value_max), classes} end)
```

<!-- livebook:{"output":true} -->

```
#Stream<[
  enum: #Function<3.112894672/2 in Task.build_stream/3>,
  funs: [#Function<50.38948127/1 in Stream.map/2>]
]>
```

```elixir
model =
  Axon.input("input", shape: {nil, image_w, image_h, image_channels})
  |> Axon.conv(64, kernel_size: {3, 3})
  |> Axon.batch_norm()
  |> Axon.relu()
  |> Axon.max_pool(kernel_size: {2, 2})
  |> Axon.conv(128, kernel_size: {3, 3})
  |> Axon.batch_norm()
  |> Axon.relu()
  |> Axon.max_pool(kernel_size: {2, 2})
  |> Axon.flatten()
  |> Axon.dense(256, activation: :relu)
  |> Axon.dropout()
  |> Axon.dense(2, activation: :softmax)
```

<!-- livebook:{"output":true} -->

```
#Axon<
  inputs: %{"input" => {nil, 256, 256, 3}}
  outputs: "softmax_0"
  nodes: 15
>
```

```elixir
optimizer = Polaris.Optimizers.adam(learning_rate: 1.0e-3)
centralized_optimizer = Polaris.Updates.compose(Polaris.Updates.centralize(), optimizer)
epochs = 4

model_state =
  model
  |> Axon.Loop.trainer(:binary_cross_entropy, centralized_optimizer, log: 1)
  |> Axon.Loop.metric(:accuracy)
  |> Axon.Loop.run(data, %{}, epochs: epochs, iterations: 30, compiler: EXLA)
```

<!-- livebook:{"output":true} -->

```

08:12:25.650 [debug] Forwarding options: [compiler: EXLA] to JIT compiler

08:12:37.790 [warning] Allocator (GPU_0_bfc) ran out of memory trying to allocate 480.50MiB (rounded to 503840768)requested by op 

08:12:37.791 [info] BFCAllocator dump for GPU_0_bfc

08:12:37.791 [info] Bin (256): 	Total Chunks: 102, Chunks in use: 102. 25.5KiB allocated for chunks. 25.5KiB in use in bin. 11.5KiB client-requested in use in bin.

08:12:37.791 [info] Bin (512): 	Total Chunks: 20, Chunks in use: 20. 10.0KiB allocated for chunks. 10.0KiB in use in bin. 10.0KiB client-requested in use in bin.

08:12:37.792 [info] Bin (1024): 	Total Chunks: 4, Chunks in use: 4. 4.0KiB allocated for chunks. 4.0KiB in use in bin. 4.0KiB client-requested in use in bin.

08:12:37.792 [info] Bin (2048): 	Total Chunks: 16, Chunks in use: 16. 44.0KiB allocated for chunks. 44.0KiB in use in bin. 43.0KiB client-requested in use in bin.

08:12:37.793 [info] Bin (4096): 	Total Chunks: 4, Chunks in use: 4. 16.0KiB allocated for chunks. 16.0KiB in use in bin. 16.0KiB client-requested in use in bin.

08:12:37.793 [info] Bin (8192): 	Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.

08:12:37.795 [info] Bin (16384): 	Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.

08:12:37.795 [info] Bin (32768): 	Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.

08:12:37.796 [info] Bin (65536): 	Total Chunks: 8, Chunks in use: 8. 624.5KiB allocated for chunks. 624.5KiB in use in bin. 576.0KiB client-requested in use in bin.

08:12:37.797 [info] Bin (131072): 	Total Chunks: 8, Chunks in use: 8. 1.20MiB allocated for chunks. 1.20MiB in use in bin. 1.12MiB client-requested in use in bin.

08:12:37.797 [info] Bin (262144): 	Total Chunks: 4, Chunks in use: 4. 1.12MiB allocated for chunks. 1.12MiB in use in bin. 1.12MiB client-requested in use in bin.

08:12:37.797 [info] Bin (524288): 	Total Chunks: 2, Chunks in use: 0. 1.54MiB allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.

08:12:37.798 [info] Bin (1048576): 	Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.

08:12:37.798 [info] Bin (2097152): 	Total Chunks: 4, Chunks in use: 1. 9.57MiB allocated for chunks. 2.50MiB in use in bin. 2.50MiB client-requested in use in bin.

08:12:37.798 [info] Bin (4194304): 	Total Chunks: 2, Chunks in use: 1. 11.25MiB allocated for chunks. 6.00MiB in use in bin. 6.00MiB client-requested in use in bin.

```

<!-- livebook:{"output":true} -->

```

08:12:37.798 [info] Bin (8388608): 	Total Chunks: 1, Chunks in use: 0. 10.49MiB allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.

08:12:37.799 [info] Bin (16777216): 	Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.

08:12:37.799 [info] Bin (33554432): 	Total Chunks: 4, Chunks in use: 4. 133.00MiB allocated for chunks. 133.00MiB in use in bin. 133.00MiB client-requested in use in bin.

08:12:37.800 [info] Bin (67108864): 	Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.

08:12:37.800 [info] Bin (134217728): 	Total Chunks: 6, Chunks in use: 4. 994.80MiB allocated for chunks. 600.00MiB in use in bin. 600.00MiB client-requested in use in bin.

08:12:37.801 [info] Bin (268435456): 	Total Chunks: 2, Chunks in use: 0. 629.22MiB allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.

08:12:37.801 [info] Bin for 480.50MiB was 256.00MiB, Chunk State: 

08:12:37.801 [info]   Size: 295.97MiB | Requested Size: 153.12MiB | in_use: 0 | bin_num: 20, prev:   Size: 223.5KiB | Requested Size: 144.0KiB | in_use: 1 | bin_num: -1, next:   Size: 144.0KiB | Requested Size: 144.0KiB | in_use: 1 | bin_num: -1

08:12:37.803 [info]   Size: 333.25MiB | Requested Size: 153.12MiB | in_use: 0 | bin_num: 20, prev:   Size: 33.25MiB | Requested Size: 33.25MiB | in_use: 1 | bin_num: -1

08:12:37.804 [info] Next region of size 1880004864

08:12:37.804 [info] Free  at 736310000000 of size 2575616 next 65

08:12:37.806 [info] InUse at 736310274d00 of size 256 next 1947

08:12:37.806 [info] InUse at 736310274e00 of size 73728 next 438

08:12:37.807 [info] InUse at 736310286e00 of size 73728 next 1507

08:12:37.807 [info] InUse at 736310298e00 of size 115712 next 1943

08:12:37.807 [info] InUse at 7363102b5200 of size 256 next 420

08:12:37.808 [info] Free  at 7363102b5300 of size 786688 next 238

08:12:37.808 [info] InUse at 736310375400 of size 256 next 259

08:12:37.809 [info] Free  at 736310375500 of size 10996736 next 3677

08:12:37.809 [info] InUse at 736310df2100 of size 2618880 next 1580

08:12:37.810 [info] Free  at 736311071700 of size 2671872 next 291

08:12:37.810 [info] InUse at 7363112fdc00 of size 147456 next 1912

08:12:37.810 [info] InUse at 736311321c00 of size 147456 next 386

08:12:37.811 [info] InUse at 736311345c00 of size 148736 next 39

08:12:37.811 [info] Free  at 73631136a100 of size 5500672 next 1277

08:12:37.811 [info] InUse at 7363118a9000 of size 294912 next 3916

08:12:37.811 [info] InUse at 7363118f1000 of size 294912 next 3549

08:12:37.812 [info] InUse at 736311939000 of size 294912 next 1748

08:12:37.812 [info] InUse at 736311981000 of size 294912 next 1904

08:12:37.812 [info] InUse at 7363119c9000 of size 73728 next 1575

08:12:37.812 [info] InUse at 7363119db000 of size 3584 next 1797

08:12:37.812 [info] InUse at 7363119dbe00 of size 3584 next 1919

08:12:37.813 [info] InUse at 7363119dcc00 of size 3584 next 230

08:12:37.813 [info] InUse at 7363119dda00 of size 3584 next 1730

08:12:37.813 [info] InUse at 7363119de800 of size 2048 next 1774

08:12:37.813 [info] InUse at 7363119df000 of size 2048 next 1639

08:12:37.813 [info] InUse at 7363119df800 of size 2048 next 1728

08:12:37.814 [info] InUse at 7363119e0000 of size 2048 next 1531

08:12:37.814 [info] InUse at 7363119e0800 of size 1024 next 1504

08:12:37.814 [info] InUse at 7363119e0c00 of size 1024 next 1498

08:12:37.815 [info] InUse at 7363119e1000 of size 1024 next 1636

08:12:37.815 [info] InUse at 7363119e1400 of size 1024 next 1742

08:12:37.815 [info] InUse at 7363119e1800 of size 512 next 1855

08:12:37.815 [info] InUse at 7363119e1a00 of size 512 next 2919

08:12:37.816 [info] InUse at 7363119e1c00 of size 512 next 342

08:12:37.816 [info] InUse at 7363119e1e00 of size 512 next 1903

08:12:37.816 [info] InUse at 7363119e2000 of size 512 next 3877

08:12:37.816 [info] InUse at 7363119e2200 of size 512 next 1808

08:12:37.816 [info] InUse at 7363119e2400 of size 512 next 1749

08:12:37.817 [info] InUse at 7363119e2600 of size 512 next 3830

08:12:37.817 [info] InUse at 7363119e2800 of size 512 next 1832

08:12:37.817 [info] InUse at 7363119e2a00 of size 512 next 1656

08:12:37.817 [info] InUse at 7363119e2c00 of size 512 next 1491

08:12:37.817 [info] InUse at 7363119e2e00 of size 512 next 3251

08:12:37.818 [info] InUse at 7363119e3000 of size 512 next 3922

08:12:37.818 [info] InUse at 7363119e3200 of size 512 next 3914

08:12:37.818 [info] InUse at 7363119e3400 of size 512 next 3755

08:12:37.818 [info] InUse at 7363119e3600 of size 512 next 1768

08:12:37.818 [info] InUse at 7363119e3800 of size 512 next 1488

08:12:37.819 [info] InUse at 7363119e3a00 of size 512 next 1831

08:12:37.819 [info] InUse at 7363119e3c00 of size 512 next 1689

08:12:37.819 [info] InUse at 7363119e3e00 of size 512 next 1846

08:12:37.819 [info] InUse at 7363119e4000 of size 256 next 1861

08:12:37.821 [info] InUse at 7363119e4100 of size 256 next 1708

08:12:37.821 [info] InUse at 7363119e4200 of size 256 next 3907

08:12:37.821 [info] InUse at 7363119e4300 of size 256 next 1828

08:12:37.822 [info] InUse at 7363119e4400 of size 256 next 1892

08:12:37.822 [info] InUse at 7363119e4500 of size 256 next 1759

08:12:37.822 [info] InUse at 7363119e4600 of size 256 next 310

08:12:37.822 [info] InUse at 7363119e4700 of size 256 next 1670

08:12:37.823 [info] InUse at 7363119e4800 of size 256 next 1758

08:12:37.823 [info] InUse at 7363119e4900 of size 256 next 3781

08:12:37.823 [info] InUse at 7363119e4a00 of size 256 next 1851

08:12:37.823 [info] InUse at 7363119e4b00 of size 256 next 1741

08:12:37.823 [info] InUse at 7363119e4c00 of size 256 next 1674

08:12:37.823 [info] InUse at 7363119e4d00 of size 256 next 274

08:12:37.824 [info] InUse at 7363119e4e00 of size 256 next 1666

08:12:37.824 [info] InUse at 7363119e4f00 of size 256 next 1879

08:12:37.824 [info] InUse at 7363119e5000 of size 256 next 1744

08:12:37.824 [info] InUse at 7363119e5100 of size 256 next 72

08:12:37.824 [info] InUse at 7363119e5200 of size 256 next 1755

08:12:37.824 [info] InUse at 7363119e5300 of size 256 next 426

08:12:37.825 [info] InUse at 7363119e5400 of size 256 next 1702

08:12:37.825 [info] InUse at 7363119e5500 of size 256 next 1785

08:12:37.825 [info] InUse at 7363119e5600 of size 256 next 1553

08:12:37.825 [info] InUse at 7363119e5700 of size 256 next 1517

08:12:37.825 [info] InUse at 7363119e5800 of size 256 next 1733

08:12:37.825 [info] InUse at 7363119e5900 of size 256 next 3883

08:12:37.825 [info] InUse at 7363119e5a00 of size 256 next 146

08:12:37.826 [info] InUse at 7363119e5b00 of size 256 next 1726

08:12:37.826 [info] InUse at 7363119e5c00 of size 256 next 1678

08:12:37.826 [info] InUse at 7363119e5d00 of size 256 next 3708

08:12:37.826 [info] InUse at 7363119e5e00 of size 256 next 1539

08:12:37.826 [info] InUse at 7363119e5f00 of size 256 next 1684

08:12:37.826 [info] InUse at 7363119e6000 of size 256 next 1751

08:12:37.827 [info] InUse at 7363119e6100 of size 256 next 1552

08:12:37.827 [info] InUse at 7363119e6200 of size 256 next 1736

08:12:37.827 [info] InUse at 7363119e6300 of size 256 next 1543

08:12:37.827 [info] InUse at 7363119e6400 of size 256 next 1723

08:12:37.827 [info] InUse at 7363119e6500 of size 256 next 325

08:12:37.827 [info] InUse at 7363119e6600 of size 256 next 1880

08:12:37.827 [info] InUse at 7363119e6700 of size 256 next 1685

08:12:37.827 [info] InUse at 7363119e6800 of size 256 next 200

08:12:37.827 [info] InUse at 7363119e6900 of size 256 next 1900

08:12:37.828 [info] InUse at 7363119e6a00 of size 256 next 1677

08:12:37.828 [info] InUse at 7363119e6b00 of size 256 next 8

08:12:37.828 [info] InUse at 7363119e6c00 of size 256 next 1767

08:12:37.828 [info] InUse at 7363119e6d00 of size 256 next 1794

08:12:37.828 [info] InUse at 7363119e6e00 of size 256 next 1790

08:12:37.828 [info] InUse at 7363119e6f00 of size 256 next 1776

08:12:37.828 [info] InUse at 7363119e7000 of size 256 next 1579

08:12:37.828 [info] InUse at 7363119e7100 of size 256 next 1550

08:12:37.829 [info] InUse at 7363119e7200 of size 256 next 1698

08:12:37.829 [info] InUse at 7363119e7300 of size 256 next 1683

08:12:37.829 [info] InUse at 7363119e7400 of size 256 next 1839

08:12:37.829 [info] InUse at 7363119e7500 of size 256 next 1556

08:12:37.829 [info] InUse at 7363119e7600 of size 256 next 56

08:12:37.830 [info] InUse at 7363119e7700 of size 256 next 1893

08:12:37.830 [info] InUse at 7363119e7800 of size 256 next 3837

08:12:37.830 [info] Free  at 7363119e7900 of size 2173440 next 2023

08:12:37.830 [info] InUse at 736311bfa300 of size 147456 next 3854

08:12:37.831 [info] InUse at 736311c1e300 of size 147456 next 3887

08:12:37.831 [info] InUse at 736311c42300 of size 228864 next 3908

08:12:37.831 [info] Free  at 736311c7a100 of size 310345216 next 347

08:12:37.832 [info] InUse at 736324471f00 of size 147456 next 199

08:12:37.832 [info] InUse at 736324495f00 of size 147456 next 447

08:12:37.832 [info] InUse at 7363244b9f00 of size 73728 next 3843

08:12:37.833 [info] InUse at 7363244cbf00 of size 73728 next 7

08:12:37.833 [info] InUse at 7363244ddf00 of size 81408 next 256

08:12:37.833 [info] InUse at 7363244f1d00 of size 256 next 1657

08:12:37.833 [info] InUse at 7363244f1e00 of size 256 next 233

08:12:37.833 [info] Free  at 7363244f1f00 of size 832512 next 3910

08:12:37.834 [info] InUse at 7363245bd300 of size 73728 next 257

08:12:37.834 [info] InUse at 7363245cf300 of size 4096 next 3790

08:12:37.834 [info] InUse at 7363245d0300 of size 4096 next 128

08:12:37.834 [info] InUse at 7363245d1300 of size 4096 next 335

08:12:37.834 [info] InUse at 7363245d2300 of size 4096 next 1908

08:12:37.834 [info] InUse at 7363245d3300 of size 3584 next 116

08:12:37.834 [info] InUse at 7363245d4100 of size 3584 next 410

08:12:37.834 [info] InUse at 7363245d4f00 of size 3584 next 307

08:12:37.834 [info] InUse at 7363245d5d00 of size 3584 next 139

08:12:37.835 [info] InUse at 7363245d6b00 of size 2048 next 1817

08:12:37.835 [info] InUse at 7363245d7300 of size 2048 next 1917

08:12:37.835 [info] InUse at 7363245d7b00 of size 2048 next 297

08:12:37.835 [info] InUse at 7363245d8300 of size 2048 next 1791

08:12:37.835 [info] InUse at 7363245d8b00 of size 256 next 1801

08:12:37.835 [info] InUse at 7363245d8c00 of size 256 next 348

08:12:37.835 [info] InUse at 7363245d8d00 of size 256 next 3741

08:12:37.835 [info] InUse at 7363245d8e00 of size 256 next 264

08:12:37.835 [info] InUse at 7363245d8f00 of size 256 next 324

08:12:37.835 [info] InUse at 7363245d9000 of size 256 next 1806

08:12:37.836 [info] InUse at 7363245d9100 of size 256 next 3713

08:12:37.836 [info] InUse at 7363245d9200 of size 256 next 117

08:12:37.836 [info] InUse at 7363245d9300 of size 256 next 1859

08:12:37.836 [info] InUse at 7363245d9400 of size 256 next 3745

08:12:37.836 [info] InUse at 7363245d9500 of size 256 next 20

08:12:37.836 [info] InUse at 7363245d9600 of size 256 next 3702

08:12:37.837 [info] InUse at 7363245d9700 of size 256 next 3879

08:12:37.837 [info] InUse at 7363245d9800 of size 256 next 226

08:12:37.837 [info] InUse at 7363245d9900 of size 256 next 1910

08:12:37.837 [info] InUse at 7363245d9a00 of size 256 next 237

08:12:37.837 [info] InUse at 7363245d9b00 of size 256 next 229

08:12:37.837 [info] InUse at 7363245d9c00 of size 256 next 136

08:12:37.837 [info] InUse at 7363245d9d00 of size 256 next 448

08:12:37.837 [info] InUse at 7363245d9e00 of size 256 next 286

08:12:37.837 [info] InUse at 7363245d9f00 of size 256 next 213

08:12:37.837 [info] InUse at 7363245da000 of size 256 next 37

08:12:37.837 [info] InUse at 7363245da100 of size 256 next 346

08:12:37.838 [info] InUse at 7363245da200 of size 256 next 165

08:12:37.838 [info] InUse at 7363245da300 of size 256 next 318

08:12:37.838 [info] InUse at 7363245da400 of size 256 next 372

08:12:37.838 [info] InUse at 7363245da500 of size 256 next 443

08:12:37.838 [info] InUse at 7363245da600 of size 256 next 193

08:12:37.838 [info] InUse at 7363245da700 of size 256 next 127

08:12:37.838 [info] InUse at 7363245da800 of size 256 next 43

08:12:37.838 [info] InUse at 7363245da900 of size 256 next 550

08:12:37.838 [info] InUse at 7363245daa00 of size 256 next 376

08:12:37.838 [info] InUse at 7363245dab00 of size 256 next 467

08:12:37.839 [info] InUse at 7363245dac00 of size 256 next 47

08:12:37.839 [info] InUse at 7363245dad00 of size 256 next 207

08:12:37.839 [info] InUse at 7363245dae00 of size 256 next 540

08:12:37.839 [info] InUse at 7363245daf00 of size 256 next 408

08:12:37.839 [info] InUse at 7363245db000 of size 256 next 279

08:12:37.839 [info] InUse at 7363245db100 of size 256 next 76

08:12:37.839 [info] InUse at 7363245db200 of size 256 next 532

08:12:37.839 [info] InUse at 7363245db300 of size 6291456 next 2880

08:12:37.840 [info] Free  at 736324bdb300 of size 150994944 next 151

08:12:37.840 [info] InUse at 73632dbdb300 of size 157286400 next 1816

08:12:37.840 [info] InUse at 7363371db300 of size 157286400 next 1783

08:12:37.840 [info] InUse at 7363407db300 of size 157286400 next 3758

08:12:37.840 [info] InUse at 736349ddb300 of size 157286400 next 1731

08:12:37.840 [info] Free  at 7363533db300 of size 262984704 next 48

08:12:37.840 [info] InUse at 736362ea8700 of size 34865152 next 319

08:12:37.840 [info] InUse at 736364fe8700 of size 34865152 next 321

08:12:37.840 [info] InUse at 736367128700 of size 34865152 next 221

08:12:37.840 [info] InUse at 736369268700 of size 34865152 next 419

08:12:37.840 [info] Free  at 73636b3a8700 of size 349442560 next 18446744073709551615

08:12:37.840 [info]      Summary of in-use Chunks by size: 

08:12:37.841 [info] 102 Chunks of size 256 totalling 25.5KiB

08:12:37.841 [info] 20 Chunks of size 512 totalling 10.0KiB

08:12:37.841 [info] 4 Chunks of size 1024 totalling 4.0KiB

08:12:37.841 [info] 8 Chunks of size 2048 totalling 16.0KiB

08:12:37.841 [info] 8 Chunks of size 3584 totalling 28.0KiB

08:12:37.841 [info] 4 Chunks of size 4096 totalling 16.0KiB

08:12:37.841 [info] 6 Chunks of size 73728 totalling 432.0KiB

08:12:37.841 [info] 1 Chunks of size 81408 totalling 79.5KiB

08:12:37.842 [info] 1 Chunks of size 115712 totalling 113.0KiB

08:12:37.842 [info] 6 Chunks of size 147456 totalling 864.0KiB

08:12:37.842 [info] 1 Chunks of size 148736 totalling 145.2KiB

08:12:37.842 [info] 1 Chunks of size 228864 totalling 223.5KiB

08:12:37.842 [info] 4 Chunks of size 294912 totalling 1.12MiB

08:12:37.842 [info] 1 Chunks of size 2618880 totalling 2.50MiB

08:12:37.842 [info] 1 Chunks of size 6291456 totalling 6.00MiB

08:12:37.842 [info] 4 Chunks of size 34865152 totalling 133.00MiB

08:12:37.842 [info] 4 Chunks of size 157286400 totalling 600.00MiB

08:12:37.842 [info] Sum Total of in-use chunks: 744.53MiB

08:12:37.843 [info] Total bytes in pool: 1880004864 memory_limit_: 1880005017 available bytes: 153 curr_region_allocation_bytes_: 3760010240

08:12:37.843 [info] Stats: 
Limit:                      1880005017
InUse:                       780699904
MaxInUse:                   1871703296
NumAllocs:                       97350
MaxAllocSize:               1152347648
Reserved:                            0
PeakReserved:                        0
LargestFreeBlock:                    0


08:12:37.843 [warning] **________________*_______***********************************____________*********__________________

08:12:37.843 [error] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 503840768 bytes.
BufferAssignment OOM Debugging.
BufferAssignment stats:
             parameter allocation:         0B
              constant allocation:       273B
        maybe_live_out allocation:    1.88GiB
     preallocated temp allocation:     7.5KiB
  preallocated temp fragmentation:     1.2KiB (16.09%)
                 total allocation:    1.88GiB
              total fragmentation:     7.8KiB (0.00%)
Peak buffers:
	Buffer 1:
		Size: 480.50MiB
		XLA Label: fusion
		Shape: f32[492032,256]
		==========================

	Buffer 2:
		Size: 480.50MiB
		XLA Label: fusion
		Shape: f32[492032,256]
		==========================

	Buffer 3:
		Size: 480.50MiB
		XLA Label: fusion
		Shape: f32[125960192]
		==========================

	Buffer 4:
		Size: 480.50MiB
		XLA Label: fusion
		Shape: f32[492032,256]
		==========================

	Buffer 5:
		Size: 288.0KiB
		XLA Label: fusion
		Shape: f32[3,3,64,128]
		==========================

	Buffer 6:
		Size: 288.0KiB
		XLA Label: fusion
		Shape: f32[3,3,64,128]
		==========================

	Buffer 7:
		Size: 288.0KiB
		XLA Label: fusion
		Shape: f32[73728]
		==========================

	Buffer 8:
		Size: 288.0KiB
		XLA Label: fusion
		Shape: f32[3,3,64,128]
		==========================

	Buffer 9:
		Size: 6.8KiB
		XLA Label: fusion
		Shape: f32[3,3,3,64]
		==========================

	Buffer 10:
		Size: 6.8KiB
		XLA Label: fusion
		Shape: f32[3,3,3,64]
		==========================

	Buffer 11:
		Size: 6.8KiB
		XLA Label: fusion
		Shape: f32[1728]
		==========================

	Buffer 12:
		Size: 6.8KiB
		XLA Label: fusion
		Shape: f32[3,3,3,64]
		==========================

	Buffer 13:
		Size: 2.0KiB
		XLA Label: fusion
		Shape: f32[256,2]
		==========================

	Buffer 14:
		Size: 2.0KiB
		XLA Label: fusion
		Shape: f32[512]
		==========================

	Buffer 15:
		Size: 2.0KiB
		XLA Label: fusion
		Shape: f32[256,2]
		==========================



```

## Testing

```elixir
test_directory = "/home/le-moski/Documents/FEFU/7S/AI/lab1-dirtyplates/platesv2/plates/test/*.jpg"

test_filenames =
  Path.wildcard(test_directory)
  |> Stream.chunk_every(batch_size, batch_size)

test_data_raw =
  test_filenames
  |> Task.async_stream(fn batch ->
    batch
    |> Enum.map(fn filename ->
      {:ok, img} = StbImage.read_file(filename, channels: image_channels)
      StbImage.resize(img, image_h, image_w) |> StbImage.to_nx()
    end)
    |> Nx.stack()
  end)

test_data =
  test_data_raw
  |> Stream.map(fn {:ok, batch} -> batch |> Nx.divide(channel_value_max) end)
```

<!-- livebook:{"output":true} -->

```
#Stream<[
  enum: #Function<3.112894672/2 in Task.build_stream/3>,
  funs: [#Function<50.38948127/1 in Stream.map/2>]
]>
```

```elixir
predictions =
  test_data
  |> Stream.map(fn batch ->
    Axon.predict(model, model_state, batch)
  end)

prediction_labels =
  predictions
  |> Stream.map(fn batch ->
    batch
    |> Nx.to_list()
    |> Enum.map(
      &if &1 |> Enum.at(0) >= &1 |> Enum.at(1) do
        "Cleaned"
      else
        "Dirty"
      end
    )
  end)
```

<!-- livebook:{"output":true} -->

```
#Stream<[
  enum: #Function<3.112894672/2 in Task.build_stream/3>,
  funs: [#Function<50.38948127/1 in Stream.map/2>, #Function<50.38948127/1 in Stream.map/2>,
   #Function<50.38948127/1 in Stream.map/2>]
]>
```

```elixir
predictions |> Enum.to_list()
```

<!-- livebook:{"output":true} -->

```
[
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186328.98827>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN]
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186328.98830>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN]
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186328.98833>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN]
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186328.98836>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN]
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186325.94801>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN]
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186325.94804>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN]
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186325.94807>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN]
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186325.94810>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN]
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186325.94813>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN]
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186328.98846>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN]
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186321.95355>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, ...]
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186321.95357>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186322.99151>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, ...],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186328.98858>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186325.94881>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, ...],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186328.98903>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186326.102539>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, ...],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186328.98921>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186329.96026>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, ...],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186325.94929>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186326.102571>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, ...],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186330.95297>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186321.95421>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, ...],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186321.95441>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186326.102666>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, ...],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186326.102702>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186329.96108>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, ...],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186329.96120>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186329.96125>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, ...],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186321.95471>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186326.102744>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, ...],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186328.99042>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186330.95355>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, ...],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186330.95367>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186325.95024>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, ...],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186330.95386>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[20][2]
    EXLA.Backend<cuda:0, 0.691480651.217186330.95409>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, ...],
      ...
    ]
  >,
  #Nx.Tensor<
    f32[4][2]
    EXLA.Backend<cuda:0, 0.691480651.217186330.95412>
    [
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN],
      [NaN, NaN]
    ]
  >
]
```

```elixir
results = Stream.zip([test_filenames, test_data_raw, prediction_labels])
```

<!-- livebook:{"output":true} -->

```
#Function<76.38948127/2 in Stream.zip_with/2>
```

```elixir
csv =
  results
  |> Stream.flat_map(fn batch ->
    {filenames_batch, _, labels_batch} = batch

    Enum.zip([filenames_batch, labels_batch])
    |> Enum.map(fn {filename, label} ->
      %{id: Path.basename(filename, ".jpg"), label: String.downcase(label)}
    end)
  end)
  |> CSV.encode(headers: [:id, :label])
  |> Enum.join()

File.write!("/home/le-moski/Documents/FEFU/7S/AI/lab1-dirtyplates/results.csv", csv)
```

<!-- livebook:{"output":true} -->

```
:ok
```

```elixir
visualization =
  results
  |> Stream.map(fn batch ->
    {filenames_batch, {:ok, raw_batch}, labels_batch} = batch

    filenames_batch
    |> Stream.with_index(fn filename, index ->
      Kino.Layout.grid(
        [
          Kino.Markdown.new("# " <> Path.basename(filename)),
          raw_batch |> Nx.to_list() |> Enum.at(index) |> Nx.tensor(type: :u8) |> Kino.Image.new(),
          labels_batch |> Enum.at(index) |> Kino.Markdown.new()
        ],
        boxed: true
      )
    end)
    |> Enum.to_list()
    |> Kino.Layout.grid()
  end)
```

<!-- livebook:{"output":true} -->

```
#Stream<[
  enum: #Function<76.38948127/2 in Stream.zip_with/2>,
  funs: [#Function<50.38948127/1 in Stream.map/2>]
]>
```

```elixir
visualization |> Enum.at(0)
```

```elixir
visualization |> Enum.at(1)
```