Skip to content

Commit

Permalink
Scala inference memory leak fix apache#11204 (apache#11216)
Browse files Browse the repository at this point in the history
* Fixes Scala memory leak (apache#10436)

* Replaced the copy and disposed of sliced ndArray to resolve memory leak

* Wrapped disposes in a finally to ensure they are called.
  • Loading branch information
andrewfayres authored and marcoabreu committed Jun 14, 2018
1 parent 2a7a2e6 commit 014338a
Showing 1 changed file with 18 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -224,13 +224,24 @@ class FeedForward private(
var i = 0
while (data.hasNext && i != numBatch) {
val batch = data.next()
i += 1
ExecutorManager.loadData(batch, dataArrays)
predExec.forward(isTrain = false)
val padded = batch.pad
val realSize = batchSize - padded
for ((list, nd) <- outputs zip predExec.outputs) {
list += nd.slice(0, realSize).copy()
try {
i += 1
ExecutorManager.loadData(batch, dataArrays)
predExec.forward(isTrain = false)
val padded = batch.pad
val realSize = batchSize - padded
for ((list, nd) <- outputs zip predExec.outputs) {
// The slice is being written to a value so that dispose can be called after the copy.
// The one liner nd.slice().copy() leads to leaking the memory of the slice.
val ndSliced = nd.slice(0, realSize)
try {
list += ndSliced.copy()
} finally {
ndSliced.dispose()
}
}
} finally {
batch.dispose()
}
}
// TODO(Yizhi): we can use Symbol.concat to do the same thing. Can it be more efficient?
Expand Down

0 comments on commit 014338a

Please sign in to comment.