Reuse executor for repeated prediction with different checkpoints
- 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.