Axon memory usage high

Elixir Framework Forums Nx Forum
,
1

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)
1 Like
2

I believe you might be missing Nx.default_backend(EXLA.Backend)

Your data occupies around 280MB, if I calculated this right. If you didn’t set the default backend, that’ll be allocated in Nx.BinaryBackend. Then, at each iteration of the loop, Nx will copy the data over to EXLA due to your choice of compiler, at least doubling the memory usage.

Then, you might be running into the GC not doing its work fast enough.

1 Like
3

Actually, the data occupies around 1GB by itself according to my calcs. I’d missed the implicit f32 conversion.

4

My notebook starts with

    Nx.Defn.default_options(compiler: EXLA)
    Nx.global_default_backend(EXLA.Backend)

I can post the whole thing, but I really see this frequently, how can I debug and know if this is a memory leak? So I can at least know it’s a real issue.

5

Ah right. I thought that was the whole Nx code.

You could add a Stream.map function that calls :erlang.gc before every epoch at least, and that will eliminate the possibility of gc being too slow.

Also, which EXLA version are you using?

1 Like
6

Everything is on latest version, exla 0.7.1.

Sorry, I don’t get that, where should I put Stream.map?

7

Put that as your Axon.Loop data input! So you’d wrap the current input Enumerable and pass that

8

Thank you @polvalente for mentioning garbage collection.

If I set the :garbage_collect option to true with Axon.Loop.run then the memory usage is low and it doesn’t increase. If I don’t set it, then it grows continuously during training.

This solves the issue.

3 Likes