I’m training simple feed forward Neural Networks on CPU and I often see the memory usage go up over 128GB of RAM.

Is this normal? The training data is 70000x4000 tensor, when this gets made I see a bump in memory usage. But then when training loop starts it keeps increasing until out of memory is reached. Should memory keep rapidly increasing during training?

My code seems simple:

      {train_data, test_data} =
        data
        |> Nx.tensor(type: :u8)
        |> Nx.divide(255.0)
        |> Nx.reshape({count, 4000})
        |> Nx.to_batched(batch_size, leftover: :discard)
        |> Enum.split(train_batches_count)

      {train_labels, test_labels} =
        series["label"]
        |> Explorer.Series.to_tensor()
        |> Nx.new_axis(-1)
        |> Nx.equal(Nx.tensor([0, 1]))
        |> Nx.to_batched(batch_size, leftover: :discard)
        |> Enum.split(train_batches_count)

      optimizer = Polaris.Optimizers.adamw(learning_rate: learning_rate)

      model_params =
        model
        |> Axon.Loop.trainer(:categorical_cross_entropy, optimizer)
        |> Axon.Loop.run(Stream.zip(train_data, train_labels), %{}, epochs: 5, compiler: EXLA)