avik-pal / Wandb.jl

Unofficial Julia bindings for logging experiments to wandb.ai
https://avik-pal.github.io/Wandb.jl/stable/
MIT License
80 stars 10 forks source link

Easier way to specify step #7

Closed bhatiaabhinav closed 2 years ago

bhatiaabhinav commented 2 years ago

In the original API, we can specify the step as a separate argument.

wandb.log({"Metric/loss": 10}, step=1000)   # loss is 10 at 1000th step.

In this Wandb.jl library, there is:

lg = WandbLogger(project="blabla", name="blablablba")
log(lg, Dict("Metric/loss"=>10), step=1000)   # loss is 10 at 1000th step.

But with the recommended logging API, there is no simple way:

with_logger(lg) do
    @info "Metric" loss=10   # But where do I specify step?
end

Internally, the step lg.global_step is incremented by lg.step_increment (=1 by default) for every @info call. In order to control the step, we need to do:

with_logger(lg) do
     lg.global_step = 999
    @info "Metric" loss=10   # Will be logged to step=1000
end

How about having a simple API like:

with_logger(lg) do
    @info "Metric" loss=10  step=1000  # More consistent with the wandb API
end
bhatiaabhinav commented 2 years ago

Before discovering this awesome library, I was using wandb in Julia using a macro I had designed, which allowed me to do:

@wandb Metric/loss=10 step=1000

The code for the above is:

using PyCall

macro wandblog(exs...)
    step = nothing
    kwargs = Any[]
    for ex in exs
        if ex isa Expr && ex.head === :(=)
            k, v = ex.args
            if !(k isa Symbol)
                k = Symbol(k)
            end
            if k == :step
                step = esc(v)
                continue
            end
            push!(kwargs, Expr(:kw, k, esc(v)))
        elseif ex isa Expr && ex.head === :... # Keyword splatting
            push!(kwargs, esc(ex))
        else
            k, v = ex, ex
            if !(k isa Symbol)
                k = Symbol(k)
            end
            if k == :step
                step = esc(v)
                continue
            end
            push!(kwargs, Expr(:kw, k, esc(v)))
        end
    end
    return quote
        kw = Dict(pairs((;$(kwargs...))))
        wandb = pyimport("wandb")
        wandb.log(kw, step=$step)
    end
end

This logs using the current wandb instance, just like python. See if you can add a similar API to this library? So that there is no need to change the logger. Something like:

@wandb "Metric" loss=10 step=1000

and an additional macro if you want to specify the wandblogger:

@wandb lg Metric/loss=10 step=1000
avik-pal commented 2 years ago
with_logger(lg) do
    @info "Metric" loss=10  step=1000  # More consistent with the wandb API
end

This should get logged as Metric/loss and Metric/step. You could change the x-axis of the plots in the web app.

I have a couple of issues with the stated changes:

  1. Changing the API as suggested means we can never log something named step.
  2. The @wandb macro doesn't seem to do anything which @info/@warn,... don't already do
bhatiaabhinav commented 2 years ago

See I believe a library called Wandb.jl should not be missing basic functionality of the original wandb API. Changing the x-axis the way you suggest is a hack, not a proper solution. When one needs to log a bunch of metrics, at different places in the code, but to the same step, (which is common in RL), it is hackish to log them to different wandb steps in favor of logging them to the same "Metric/step". How about reserving a keyword say _step or wandb_step to specify the wandb step?

avik-pal commented 2 years ago

The macro API using @info/@warn is meant to be used when the user wants to combine different loggers. (NOTE: This is not the recommended API. If I have mentioned it anywhere in the docs I should drop that)

Additionally special-casing the step keyword means @info will log differently to the different loggers combined via https://github.com/JuliaLogging/LoggingExtras.jl.

I am not necessarily against adding the @wandb macro (here we could reserve a kwarg for step). But it would need to account for which kind of logger is currently being used -- WandbLogger, WandbLoggerMPI, WandbBackend, etc. If you can create a PR with that I would go ahead and review it.