Softmax refactor #7
Conversation
This commit may add some overhead of managing NDArray for each fallback.
Conflicts: src/operator/nn/mkldnn/mkldnn_batch_norm-inl.h
2. Add memory into signature; 3. Try to split BatchNorm into .h file and .cc file. Will finish it after backward code is refactored.
Caching primitive for BatchNorm forward computation
| const OpReqType &req) { | ||
| // mkldnn::softmax_forward::primitive_desc | ||
| auto input_mem = in_data.GetMKLDNNData(); | ||
| auto output_mem = out_data.GetMKLDNNData(); |
There was a problem hiding this comment.
Calling GetMKLDNNData() will create memory for out_data with default layout. But out_data may need a specific layout which is same as in_data?
There was a problem hiding this comment.
correct. This is a problem in my original code.
|
|
||
| mkldnn::memory::primitive_desc data_mpd = input_mem->get_primitive_desc(); | ||
| mkldnn::memory::desc data_md = data_mpd.desc(); | ||
| auto cpu_engine = data_mpd.get_engine(); |
| bool is_train, | ||
| const NDArray &in_data, | ||
| const NDArray &out_data, | ||
| const OpReqType &req) { |
| ~MKLDNNSoftmaxFwd() {} | ||
| void SetDataHandle(const NDArray &in_data, | ||
| const NDArray &out_data, | ||
| const OpReqType &req) { |
There was a problem hiding this comment.
yes, I will remove this input.
…em; 3.cpu engine use encapsulated func
| const OpReqType &req) { | ||
| // mkldnn::softmax_forward::primitive_desc | ||
| auto input_mem = in_data.GetMKLDNNData(); | ||
| auto output_mem = out_data.GetMKLDNNData(); |
There was a problem hiding this comment.
correct. This is a problem in my original code.
src/operator/nn/softmax.cc
Outdated
| if (dev_mask == mshadow::cpu::kDevMask) | ||
| // We only run MKLDNN op if it runs on CPU and the input data is MKLDNN | ||
| // format. | ||
| if (dev_mask == mshadow::cpu::kDevMask && (*in_attrs)[0] == kMKLDNNStorage) |
There was a problem hiding this comment.
Remove kMKLDNNStorage. There isn't kMKLDNNStorage any more.
| MKLDNNSmSignature key(param); | ||
| key.AddSign(ctx.is_train); | ||
| key.AddSign(param.axis); | ||
| key.AddSign(in_data); |
There was a problem hiding this comment.
We should add MKLDMM memory as part of a signature. We need layout info.
| mkldnn::memory::primitive_desc data_mpd = input_mem->get_primitive_desc(); | ||
| auto output_mem = CreateMKLDNNMem(out_data, data_mpd, req).second; |
There was a problem hiding this comment.
if you use CreateMKLDNNMem, you have to use it with CommitOutput. You can't use CreateMKLDNNMem alone.
You could use the CreateMKLDNNData method in NDArray, as batchnorm does.
|
softmax has its unittest already, and mkldnn can pass: nosetests tests/python/unittest/test_operator.py:test_new_softmax |
zheng-da
force-pushed
the
refactor
branch
8 times, most recently
from
ba474be
to
5cb7ca0
Compare
Description
Checklist
Essentials
make lint)Changes
Comments