fix lstm bug when hidden_size is zero

This commit is contained in:
baihuawei 2020-06-17 09:54:31 +08:00
parent 30b19165a2
commit 7f7c006acf
3 changed files with 13 additions and 0 deletions

View File

@ -22,6 +22,10 @@
namespace mindspore {
namespace kernel {
void LstmCPUKernel::InitKernel(const CNodePtr &kernel_node) {
#ifdef PLATFORM_86
_MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_ON);
_MM_SET_DENORMALS_ZERO_MODE(_MM_DENORMALS_ZERO_ON);
#endif
MS_EXCEPTION_IF_NULL(kernel_node);
using tag = dnnl::memory::format_tag;
using dim = dnnl::memory::dims;

View File

@ -16,6 +16,12 @@
#ifndef MINDSPORE_CCSRC_KERNEL_CPU_LSTM_CPU_KERNEL_H_
#define MINDSPORE_CCSRC_KERNEL_CPU_LSTM_CPU_KERNEL_H_
#if defined(__x86_64__) || defined(__amd64__) || defined(_M_IX86) || defined(_M_X64)
#define PLATFORM_86
#endif
#ifdef PLATFORM_86
#include <pmmintrin.h>
#endif
#include <vector>
#include <memory>
#include "kernel/cpu/mkldnn/mkl_cpu_kernel.h"

View File

@ -23,6 +23,7 @@ from mindspore.common.parameter import Parameter, ParameterTuple
from mindspore.common.tensor import Tensor
from mindspore.nn.cell import Cell
from mindspore.ops import operations as P
from ..._checkparam import Rel
__all__ = ['LSTM', 'LSTMCell']
@ -123,6 +124,8 @@ class LSTM(Cell):
self.num_layers = num_layers
self.has_bias = has_bias
self.batch_first = validator.check_value_type("batch_first", batch_first, [bool], self.cls_name)
self.hidden_size = validator.check_integer("hidden_size", hidden_size, 0, Rel.GT, self.cls_name)
self.num_layers = validator.check_integer("num_layers", num_layers, 0, Rel.GT, self.cls_name)
self.dropout = float(dropout)
self.bidirectional = bidirectional
if self.batch_first: