Logging plots causes memory leak

[gabrevaya]

Hi, I've just realized that the cause of the memory leak that I was facing was due to logging (maybe in the wrong way) with Wandb.jl

I was using it like this:

function train_with_Wandb(dict)
    lg = WandbLogger(
                    project = "Testing",
                    name = "$(dict[:name])-$(now())",
                    config = dict,
                    )

    with_logger(lg) do
        try
            train(dict, lg)
        finally
            Wandb.wandb.finish()
        end
    end
    close(lg)
end

In order to be able to pass lg to the train function, and there to another training_loop! where I logged:

Wandb.log(
    lg,
    Dict(
        "Training/Loss" => loss,
        "Testing/ValLoss" => val_loss,
        "LatentPlots" => wandb_plot,
    ),
    step = epoch,
)

where wandb_plot is:

wandb_plot, axs = plot_for_wandb(cpu(z_j), cpu(ẑ_j))

function plot_for_wandb(z, ẑ)

    plt = pyimport("matplotlib.pyplot")
    N = Int(size(z, 1)/2)
    fig, axs = plt.subplots(2, N)
    for i in 1:N
        axs[1,i].plot(z[i,:], label = "true")
        axs[1,i].plot(ẑ[i,:], label = "model")

        axs[2,i].plot(z[i+N,:], label = "true")
        axs[2,i].plot(ẑ[i+N,:], label = "predicted")

        axs[2,i].set_xlabel("time steps")
        axs[i].grid(true)
        axs[i+N].grid(true)

        axs[i+N].label_outer()
        axs[i].label_outer()
    end
    plt.close()
    return fig, axs
end

Am I doing something wrong here? Can you see the reason for the memory leak? If I log the losses but not the plots, I don't see the memory leak.

[gabrevaya]

I'm sorry, I've just noticed that the plotting with matplotlib alone is already causing a smaller memory leak which is probably potentiated when logged by Wandb. Maybe there is also a memory leak caused by the way I'm using Wandb so I'll wait until any opinions before closing the issue.

[avik-pal]

Seems like you resolved the issue. Also, from my recent experiments with CairoMakie with Wandb, I don't seem to find any issues. Feel free to reopen if this shows up again.