This is an automated email from the ASF dual-hosted git repository.
patriczhao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/master by this push:
new 8b2ae57 Fix MKLDNNDataReorderAsync (#15163)
8b2ae57 is described below
commit 8b2ae571964ddebfe3965ae109364f02d10812e6
Author: Zhennan Qin <[email protected]>
AuthorDate: Tue Jun 11 00:53:17 2019 -0500
Fix MKLDNNDataReorderAsync (#15163)
---
src/engine/naive_engine.cc | 8 ++++----
src/ndarray/ndarray.cc | 15 ++++++++++-----
tests/python/mkl/test_mkldnn.py | 28 +++++++++++++++++++++++-----
3 files changed, 37 insertions(+), 14 deletions(-)
diff --git a/src/engine/naive_engine.cc b/src/engine/naive_engine.cc
index 7291f46..c321e48 100644
--- a/src/engine/naive_engine.cc
+++ b/src/engine/naive_engine.cc
@@ -171,10 +171,6 @@ class NaiveEngine final : public Engine {
opr->opr_profile.reset(new profiler::ProfileOperator(opr->opr_name,
attrs.release()));
opr->opr_profile->start(exec_ctx.dev_type, exec_ctx.dev_id);
}
- // increment mutable var version
- for (auto var : mutable_vars) {
- ++var->version_;
- }
if (exec_ctx.dev_mask() == gpu::kDevMask) {
#if MXNET_USE_CUDA
size_t dev_id = static_cast<size_t>(exec_ctx.dev_id);
@@ -194,6 +190,10 @@ class NaiveEngine final : public Engine {
} else {
exec_fun(RunContext{exec_ctx, &cpu_stream_, nullptr, false}, callback);
}
+ // increment mutable var version
+ for (auto var : mutable_vars) {
+ ++var->version_;
+ }
CHECK(this->req_completed_)
<< "NaiveEngine only support synchronize Push so far";
if (profiling) {
diff --git a/src/ndarray/ndarray.cc b/src/ndarray/ndarray.cc
index af9e7a7..bee8bef 100644
--- a/src/ndarray/ndarray.cc
+++ b/src/ndarray/ndarray.cc
@@ -612,12 +612,17 @@ void NDArray::MKLDNNDataReorderAsync(const
mkldnn::memory::primitive_desc &desc)
std::vector<Engine::VarHandle> const_vars;
std::vector<Engine::VarHandle> mutable_vars(1, this->var());
NDArray tmp = *this;
+ const auto version = this->version();
Engine::Get()->PushAsync(
- [tmp, desc](RunContext ctx, Engine::CallbackOnComplete on_complete) {
- tmp.ptr_->MKLDNNDataReorder(desc);
- on_complete();
- }, ctx(), const_vars, mutable_vars,
- FnProperty::kNormal, 0, "Reorder");
+ [tmp, version, desc](RunContext ctx, Engine::CallbackOnComplete
on_complete) {
+ // MXNet will try to reuse NDArray from memory planning, so we need to
ensure
+ // the NDArray is still holding the original trunk data.
+ if (tmp.version() == version) {
+ tmp.ptr_->MKLDNNDataReorder(desc);
+ }
+ on_complete();
+ },
+ ctx(), const_vars, mutable_vars, FnProperty::kNormal, 0, "Reorder");
}
const mkldnn::memory *NDArray::GetMKLDNNData() const {
diff --git a/tests/python/mkl/test_mkldnn.py b/tests/python/mkl/test_mkldnn.py
index b4cfde2..662edcf 100644
--- a/tests/python/mkl/test_mkldnn.py
+++ b/tests/python/mkl/test_mkldnn.py
@@ -24,6 +24,7 @@ import numpy as np
import mxnet as mx
import unittest
from mxnet.test_utils import rand_ndarray, assert_almost_equal
+from mxnet.module import Module
from mxnet import gluon
from mxnet.gluon import nn
from mxnet.test_utils import *
@@ -496,7 +497,7 @@ def test_reshape_transpose_6d():
def __init__(self, factor):
super(Reshape2D, self).__init__()
self._factors = (int(factor),) * 2
-
+
def hybrid_forward(self, F, x):
f1, f2 = self._factors
# (N, f1*f2*C, H, W)
@@ -505,20 +506,20 @@ def test_reshape_transpose_6d():
x = F.transpose(x, (0, 1, 4, 2, 5, 3)) # (N, C, H, f1, W,
f2)
x = F.reshape(x, (0, 0, -3, -3)) # (N, C, H*f1, W*f2)
return x
-
-
+
+
class Net(gluon.HybridBlock):
def __init__(self, **kwargs):
super(Net, self).__init__(**kwargs)
with self.name_scope():
self.conv1 = nn.Conv2D(8, kernel_size=5)
self.reshape2D = Reshape2D(2)
-
+
def hybrid_forward(self, F, x):
x = self.conv1(x)
x = self.reshape2D(x)
return x
-
+
net = Net()
net.initialize(mx.init.Xavier(), ctx=mx.cpu())
net.hybridize()
@@ -526,5 +527,22 @@ def test_reshape_transpose_6d():
output = net(data)
a = output.asnumpy()
+@with_seed()
+def test_weight_async_reorder():
+ data = mx.sym.Variable("data")
+ w1 = mx.sym.Variable("1_weight")
+ w2 = mx.sym.Variable("2_weight")
+ conv1 = mx.sym.Convolution(data=data, weight=w1 + w1, num_filter=32,
no_bias=True, kernel=(3, 3))
+ conv2 = mx.sym.Convolution(data=conv1, weight=w2 + w2, num_filter=32,
no_bias=True, kernel=(1, 1))
+ mod = Module(symbol=conv2, label_names=None, context=mx.current_context())
+ mod.bind(for_training=False, data_shapes=[('data', (10, 16, 50, 50))])
+ mod.init_params(initializer=mx.init.Xavier(magnitude=2.))
+ data = [mx.random.uniform(-1.0, 1.0, shape=(10, 16, 50, 50),
ctx=mx.current_context())]
+ batch=mx.io.DataBatch(data, [])
+ for i in range(2):
+ mod.forward(batch, is_train=False)
+ for output in mod.get_outputs():
+ output.wait_to_read()
+
if __name__ == '__main__':
install.test_mkldnn_install()
