[numpy] Fix d2l performance regression (apache#15173)

* Add np array adapter decorator for layers

* Fix performance regression caused by too many conversions between nd.NDArray and np.ndarray

* Fix pylint

* Fix test backward compatibility issue

* Fix test_lambda

python/mxnet/gluon/data/vision/transforms.py

@@ -23,7 +23,7 @@
 from ...nn import Sequential, HybridSequential
 from .... import image
 from ....base import numeric_types
-from ....util import is_np_array
+from ...utils import _adapt_np_array
 
 
 class Compose(Sequential):
@@ -134,11 +134,9 @@ class ToTensor(HybridBlock):
  def __init__(self):
  super(ToTensor, self).__init__()
 
+ @_adapt_np_array
  def hybrid_forward(self, F, x):
- if is_np_array():
- x = x.as_classic_ndarray()
- out = F.image.to_tensor(x)
- return out.as_np_ndarray() if is_np_array() else out
+ return F.image.to_tensor(x)
 
 
 class Normalize(HybridBlock):

python/mxnet/gluon/loss.py

@@ -29,7 +29,8 @@
 from .. import ndarray
 from ..base import numeric_types
 from .block import HybridBlock
-from .utils import _to_classic_arrays, _to_np_arrays
+from .utils import _adapt_np_array
+from ..util import is_np_array
 
 
 def _apply_weighting(F, loss, weight=None, sample_weight=None):
@@ -54,7 +55,10 @@ def _apply_weighting(F, loss, weight=None, sample_weight=None):
  Weighted loss
  """
  if sample_weight is not None:
- loss = F.broadcast_mul(loss, sample_weight)
+ if is_np_array():
+ loss = loss * sample_weight
+ else:
+ loss = F.broadcast_mul(loss, sample_weight)
 
  if weight is not None:
  assert isinstance(weight, numeric_types), "weight must be a number"
@@ -65,7 +69,11 @@ def _apply_weighting(F, loss, weight=None, sample_weight=None):
 
 def _reshape_like(F, x, y):
  """Reshapes x to the same shape as y."""
- return x.reshape(y.shape) if F is ndarray else F.reshape_like(x, y)
+ if F is ndarray:
+ return x.reshape(y.shape)
+ elif is_np_array():
+ F = F.npx
+ return F.reshape_like(x, y)
 
 
 class Loss(HybridBlock):
@@ -136,14 +144,16 @@ def __init__(self, weight=1., batch_axis=0, **kwargs):
  super(L2Loss, self).__init__(weight, batch_axis, **kwargs)
 
  def hybrid_forward(self, F, pred, label, sample_weight=None):
- # TODO(junwu): This is a temp solution to reuse legacy ops for np.ndarray.
- # We should rewrite this with np/npx ops.
- pred, label, sample_weight = _to_classic_arrays(pred, label, sample_weight)
  label = _reshape_like(F, label, pred)
- loss = F.square(label - pred)
+ loss = F.np.square(label - pred) if is_np_array() else F.square(label - pred)
  loss = _apply_weighting(F, loss, self._weight / 2, sample_weight)
- out = F.mean(loss, axis=self._batch_axis, exclude=True)
- return _to_np_arrays(out)
+ if is_np_array():
+ if F is ndarray:
+ return F.np.mean(loss, axis=tuple(range(1, loss.ndim)))
+ else:
+ return F.npx.batch_flatten(loss).mean(axis=1)
+ else:
+ return F.mean(loss, axis=self._batch_axis, exclude=True)
 
 
 class L1Loss(Loss):
@@ -178,15 +188,12 @@ class L1Loss(Loss):
  def __init__(self, weight=None, batch_axis=0, **kwargs):
  super(L1Loss, self).__init__(weight, batch_axis, **kwargs)
 
+ @_adapt_np_array
  def hybrid_forward(self, F, pred, label, sample_weight=None):
- # TODO(junwu): This is a temp solution to reuse legacy ops for np.ndarray.
- # We should rewrite this with np/npx ops.
- pred, label, sample_weight = _to_classic_arrays(pred, label, sample_weight)
  label = _reshape_like(F, label, pred)
  loss = F.abs(label - pred)
  loss = _apply_weighting(F, loss, self._weight, sample_weight)
- out = F.mean(loss, axis=self._batch_axis, exclude=True)
- return _to_np_arrays(out)
+ return F.mean(loss, axis=self._batch_axis, exclude=True)
 
 
 class SigmoidBinaryCrossEntropyLoss(Loss):
@@ -251,11 +258,8 @@ def __init__(self, from_sigmoid=False, weight=None, batch_axis=0, **kwargs):
  weight, batch_axis, **kwargs)
  self._from_sigmoid = from_sigmoid
 
+ @_adapt_np_array
  def hybrid_forward(self, F, pred, label, sample_weight=None, pos_weight=None):
- # TODO(junwu): This is a temp solution to reuse legacy ops for np.ndarray.
- # We should rewrite this with np/npx ops.
- pred, label, sample_weight, pos_weight =\
- _to_classic_arrays(pred, label, sample_weight, pos_weight)
  label = _reshape_like(F, label, pred)
  if not self._from_sigmoid:
  if pos_weight is None:
@@ -277,8 +281,7 @@ def hybrid_forward(self, F, pred, label, sample_weight=None, pos_weight=None):
  loss = -(F.broadcast_mul(F.log(pred + eps) * label, pos_weight)
  + F.log(1. - pred + eps) * (1. - label))
  loss = _apply_weighting(F, loss, self._weight, sample_weight)
- out = F.mean(loss, axis=self._batch_axis, exclude=True)
- return _to_np_arrays(out)
+ return F.mean(loss, axis=self._batch_axis, exclude=True)
 
 
 SigmoidBCELoss = SigmoidBinaryCrossEntropyLoss
@@ -354,10 +357,8 @@ def __init__(self, axis=-1, sparse_label=True, from_logits=False, weight=None,
  self._sparse_label = sparse_label
  self._from_logits = from_logits
 
+ @_adapt_np_array
  def hybrid_forward(self, F, pred, label, sample_weight=None):
- # TODO(junwu): This is a temp solution to reuse legacy ops for np.ndarray.
- # We should rewrite this with np/npx ops.
- pred, label = _to_classic_arrays(pred, label)
  if not self._from_logits:
  pred = F.log_softmax(pred, self._axis)
  if self._sparse_label:
@@ -366,8 +367,7 @@ def hybrid_forward(self, F, pred, label, sample_weight=None):
  label = _reshape_like(F, label, pred)
  loss = -F.sum(pred * label, axis=self._axis, keepdims=True)
  loss = _apply_weighting(F, loss, self._weight, sample_weight)
- out = F.mean(loss, axis=self._batch_axis, exclude=True)
- return _to_np_arrays(out)
+ return F.mean(loss, axis=self._batch_axis, exclude=True)
 
 
 SoftmaxCELoss = SoftmaxCrossEntropyLoss

python/mxnet/gluon/nn/activations.py

@@ -22,7 +22,7 @@
 
 from ... import initializer
 from ..block import HybridBlock
-from ..utils import _to_classic_arrays, _to_np_arrays
+from ...util import is_np_array
 
 
 class Activation(HybridBlock):
@@ -49,9 +49,9 @@ def _alias(self):
  return self._act_type
 
  def hybrid_forward(self, F, x):
- x = _to_classic_arrays(x)
- out = F.Activation(x, act_type=self._act_type, name='fwd')
- return _to_np_arrays(out)
+ if is_np_array():
+  F = F.npx
+ return F.Activation(x, act_type=self._act_type, name='fwd')
 
  def __repr__(self):
  s = '{name}({_act_type})'

python/mxnet/gluon/nn/basic_layers.py

@@ -26,8 +26,9 @@
 
 from .activations import Activation
 from ..block import Block, HybridBlock
-from ..utils import _indent, _to_classic_arrays, _to_np_arrays
+from ..utils import _indent, _adapt_np_array
 from ... import nd, sym
+from ...util import is_np_array
 
 
 class Sequential(Block):
@@ -218,14 +219,13 @@ def __init__(self, units, activation=None, use_bias=True, flatten=True,
  self.act = None
 
  def hybrid_forward(self, F, x, weight, bias=None):
- # TODO(junwu): This is a temp solution to reuse legacy ops for np.ndarray.
- # We should rewrite this with np/npx ops.
- x, weight, bias = _to_classic_arrays(x, weight, bias)
+ if is_np_array():
+ F = F.npx
  act = F.FullyConnected(x, weight, bias, no_bias=bias is None, num_hidden=self._units,
  flatten=self._flatten, name='fwd')
  if self.act is not None:
  act = self.act(act)
- return _to_np_arrays(act)
+ return act
 
  def __repr__(self):
  s = '{name}({layout}, {act})'
@@ -265,13 +265,12 @@ def __init__(self, rate, axes=(), **kwargs):
  self._rate = rate
  self._axes = axes
 
+ @_adapt_np_array
  def hybrid_forward(self, F, x):
- x = _to_classic_arrays(x)
  if self._rate > 0:
- out = F.Dropout(x, p=self._rate, axes=self._axes, name='fwd', cudnn_off=False)
+ return F.Dropout(x, p=self._rate, axes=self._axes, name='fwd', cudnn_off=False)
  else:
- out = F.identity(x)
- return _to_np_arrays(out)
+ return F.identity(x)
 
  def __repr__(self):
  s = '{name}(p = {_rate}, axes={_axes})'
@@ -361,6 +360,7 @@ def cast(self, dtype):
  dtype = 'float32'
  super(BatchNorm, self).cast(dtype)
 
+ @_adapt_np_array
  def hybrid_forward(self, F, x, gamma, beta, running_mean, running_var):
  return F.BatchNorm(x, gamma, beta, running_mean, running_var,
  name='fwd', **self._kwargs)
@@ -414,6 +414,7 @@ def __init__(self, input_dim, output_dim, dtype='float32',
  init=weight_initializer, dtype=dtype,
  allow_deferred_init=True, grad_stype=grad_stype)
 
+ @_adapt_np_array
  def hybrid_forward(self, F, x, weight):
  return F.Embedding(x, weight, name='fwd', **self._kwargs)
 
@@ -435,6 +436,7 @@ class Flatten(HybridBlock):
  def __init__(self, **kwargs):
  super(Flatten, self).__init__(**kwargs)
 
+ @_adapt_np_array
  def hybrid_forward(self, F, x):
  return F.Flatten(x)
 
@@ -520,6 +522,7 @@ def __init__(self, axis=1, epsilon=1e-5, center=True, scale=False,
  shape=(in_channels,), init=beta_initializer,
  allow_deferred_init=True)
 
+ @_adapt_np_array
  def hybrid_forward(self, F, x, gamma, beta):
  if self._axis == 1:
  return F.InstanceNorm(x, gamma, beta,
@@ -608,6 +611,7 @@ def __init__(self, axis=-1, epsilon=1e-5, center=True, scale=True,
  shape=(in_channels,), init=beta_initializer,
  allow_deferred_init=True)
 
+ @_adapt_np_array
  def hybrid_forward(self, F, data, gamma, beta):
  norm_data = F.LayerNorm(data, gamma=gamma, beta=beta, axis=self._axis, eps=self._epsilon)
  return norm_data
@@ -792,6 +796,7 @@ def __init__(self, function, prefix=None):
  "Unrecognized function in lambda: {} of type {}"
  .format(function, type(function)))
 
+ @_adapt_np_array
  def hybrid_forward(self, F, x, *args):
  return self._func(F, x, *args)
 

python/mxnet/gluon/utils.py

@@ -38,7 +38,7 @@ class requests_failed_to_import(object):
 import numpy as np
 
 from .. import ndarray
-from ..util import is_np_shape, is_np_array
+from ..util import is_np_shape, is_np_array, wraps_safely
 
 
 def split_data(data, num_slice, batch_axis=0, even_split=True):
@@ -459,7 +459,7 @@ def _check_same_symbol_type(symbols):
  'symbols in the list to numpy symbols by calling `as_np_ndarray()` '
  'on each of them; if you want classic ndarray output(s) from the '
  'computation graph, please convert all the numpy symbols in the list '
- 'to classic symbols by calling `as_classic_ndarray()` on each of them.')
+ 'to classic symbols by calling `as_nd_ndarray()` on each of them.')
  return np_symbol if is_np_sym else classic_symbol
 
 
@@ -474,16 +474,24 @@ def _check_all_np_ndarrays(out):
  '{}'.format(str(type(array))))
 
 
-def _to_classic_arrays(*args):
+def _to_classic_arrays(*args, **kwargs):
  """Convert arrays to classic arrays. This is used in a Gluon layer for converting
  inputs of np arrays to classic arrays so that the layer built with legacy ops can still
  be used in np_array semantics."""
+ from ..numpy import ndarray as np_ndarray
+ from ..symbol.numpy import _Symbol as np_symbol
  num_inputs = len(args)
  assert num_inputs != 0
  if not is_np_array():
- return args[0] if num_inputs == 1 else args
- in_arrs = [arr if arr is None else arr.as_classic_ndarray() for arr in args]
- return in_arrs[0] if num_inputs == 1 else in_arrs
+ return args, kwargs
+ in_arrs = [arr if arr is None else arr.as_nd_ndarray() for arr in args]
+ new_kwargs = {}
+ for k, v in kwargs.items():
+ if isinstance(v, (np_ndarray, np_symbol)):
+ new_kwargs[k] = v.as_nd_ndarray()
+ else:
+ new_kwargs[k] = v
+ return in_arrs, new_kwargs
 
 
 def _to_np_arrays(*args):
@@ -496,3 +504,21 @@ def _to_np_arrays(*args):
  return args[0] if num_outputs == 1 else args
  out = [arr.as_np_ndarray() for arr in args]
  return out[0] if num_outputs == 1 else out
+
+
+# TODO(junwu): This is a temp solution for allowing basic layers
+# implemented using legacy ops to accept np.ndarrays as inputs and return
+# np.ndarrays as outputs. We should remove it after changing all the layers
+# to use np ops in np_array semantics in the future.
+def _adapt_np_array(func):
+ @wraps_safely(func)
+ def _with_np_array(*args, **kwargs):
+ assert len(args) > 2, "expect at least three arguments in args"
+ if is_np_array():
+ input_args, kwargs = _to_classic_arrays(*args[2:], **kwargs)
+ input_args = list(args[0:2]) + input_args
+ out = func(*input_args, **kwargs)
+ return _to_np_arrays(out)
+ else:
+ return func(*args, **kwargs)
+ return _with_np_array

python/mxnet/ndarray/ndarray.py

@@ -196,7 +196,7 @@ def as_np_ndarray(self):
  check_call(_LIB.MXShallowCopyNDArray(self.handle, ctypes.byref(hdl)))
  return ndarray(handle=hdl, writable=self.writable)
 
- def as_classic_ndarray(self):
+ def as_nd_ndarray(self):
  """A convenience function for creating a classic ndarray from the current
  ndarray with zero copy. For this class, it just returns itself since it is
  already a classic ndarray."""
@@ -962,7 +962,7 @@ def _at(self, idx):
  % (idx-length, length))
  check_call(_LIB.MXNDArrayAt(
  self.handle, mx_uint(idx), ctypes.byref(handle)))
- return NDArray(handle=handle, writable=self.writable)
+ return self.__class__(handle=handle, writable=self.writable)
 
  def reshape(self, *shape, **kwargs):
  """Returns a **view** of this array with a new shape without altering any data.