dmlc / dmlc/MXNet.jl

Automated adaption of batch_size to available GPU memory

Open
#83 1 comment 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

I am running a different set of models upon different input sizes and so it is not very straightforward to choose an optimal batch-size for each model/input combination.

I came up with the following solution I maybe MXNet should provide something similar.

``` julia
global const GPU_MEM_MAX = 11000 # M6000 has 12000 MiB

"""
size dim1...dim3 example (256, 256, 1)
"""
function memoryPerBatch(net, size)
exec = mx.simple_bind(net, mx.gpu(), data=(size..., 1))
dbg_str = mx.debug_str(exec)
mem_str = split(dbg_str, ['\n'])[end-2]
m = match(r"Total (\d+) MB allocated", mem_str)
finalize(exec.handle)
exec = nothing
parse(Int, m.captures[1])
end

function guessBatchSize(net, size, max = GPU_MEM_MAX)
mem = memoryPerBatch(net, size)
info("Memory required per batch: $mem MB")
nbatch = max ÷ mem
nbatch = ispow2(nbatch) ? nbatch : prevpow2(nbatch)
info("Using a batch_size of $nbatch. Total memory ≈ $(nbatch*mem) MB")
return nbatch
end
```

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the Julia prototype's memoryPerBatch and guessBatchSize functions, then trace the corresponding mx.simple_bind, mx.gpu, and mx.debug_str entry points. Define how automatic batch sizing should use available GPU memory and what result the library should return for varying models and input sizes.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.