dmlc / dmlc/MXNet.jl

Reuse executor for repeated prediction with different checkpoints

Open
#84 5 comments 0 reactions 0 assignees View on GitHub
Dominant language
No language data
Stars
369
Forks
63
PR merge metrics
No merged PRs in 30d

Description

Testing the different checkpoints of a training run requires loading checkpoints and predictions runs in a tight loop.

``` julia
net = mx.load(archfile, mx.SymbolicNode)
arch = mx.FeedForward(net, context=mx.gpu())

data = ...

for wfile in weights
saved_dict = mx.load(wfile, mx.NDArray)
arg_params = Dict{Base.Symbol, mx.NDArray}()
aux_params = Dict{Base.Symbol, mx.NDArray}()
for (k,v) in saved_dict
tp, name = split(string(k), ':')
name = symbol(name)
if tp == "arg"
arg_params[name] = v
else
aux_params[name] = v
end
end

arch.arg_params = arg_params
arch.aux_params = aux_params

pred = mx.predict(arch, data)
# For memory reclaim, eagerly finalize self.pred_exec.handle
finalize(arch.pred_exec.handle)
arch.pred_exec=nothing
end
```

Without these final two lines this easily runs out of memory for big models/batch_sizes, because the executor is not gc'd yet and we are creating a new one asking for more memory.

If we could reuse the previous executor that problem would be alleviated.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.