Fix #17164 symbolblock with BatchNorm inside during cast to fp16

python/mxnet/gluon/block.py

@@ -16,7 +16,7 @@
 # under the License.
 
 # coding: utf-8
-# pylint: disable= arguments-differ, too-many-lines
+# pylint: disable= arguments-differ, too-many-lines, reimported
 """Base container class for all neural network models."""
 __all__ = ['Block', 'HybridBlock', 'SymbolBlock']
 
@@ -25,6 +25,7 @@
 import warnings
 import re
 from collections import OrderedDict, defaultdict
+import numpy as np
 
 from ..base import mx_real_t, MXNetError
 from .. import symbol, ndarray, initializer, np_symbol
@@ -1353,6 +1354,28 @@ def _clear_cached_op(self):
  def cast(self, dtype):
  self._clear_cached_op()
  super(SymbolBlock, self).cast(dtype)
+ if np.dtype(dtype).name == 'float16':
+ # correct BatchNorm types back to float32 due to its special requirement
+ out = self._cached_graph[1]
+ params_list = out.get_internals().list_inputs()
+ for node in params_list:
+ if node.endswith('running_var'):
+ prefix = node[:-11]
+ sibs = [prefix + t for t in ('running_mean', 'gamma', 'beta')]
+ is_bn = all(p in params_list for p in sibs)
+ if is_bn:
+ self.params.get(node).cast('float32')
+ for sib in sibs:
+ self.params.get(sib).cast('float32')
+ if node.endswith('moving_var'):
+ # another convention used
+ prefix = node[:-10]
+ sibs = [prefix + t for t in ('moving_mean', 'gamma', 'beta')]
+ is_bn = all(p in params_list for p in sibs)
+ if is_bn:
+ self.params.get(node).cast('float32')
+ for sib in sibs:
+ self.params.get(sib).cast('float32')
 
  def hybrid_forward(self, F, x, *args, **kwargs):
  raise NotImplementedError

tests/python/gpu/test_gluon_gpu.py

@@ -594,6 +594,24 @@ def hybrid_forward(self, F, a, b):
  assert_raises(ValueError, lambda: foo_hybrid(mx.nd.ones((10,), ctx=mx.gpu()),
  mx.nd.ones((10,), ctx=mx.cpu())))
 
+@with_seed()
+def test_symbol_block_symbolic_bn_fp16_cast():
+ with mx.gpu(0):
+ net = mx.gluon.nn.HybridSequential()
+ sym = mx.sym.var('data')
+ conv = mx.sym.Convolution(sym, kernel=(3, 3), num_filter=16)
+ bn = mx.sym.BatchNorm(conv, name='bn_test')
+ internals = bn.get_internals()
+ net.add(mx.gluon.nn.SymbolBlock([internals['bn_test_output']], [mx.sym.var('data')]))
+ net.add(mx.gluon.nn.Conv2D(10, kernel_size=1))
+ net.initialize()
+ x = mx.nd.zeros((1, 3, 32, 32), dtype='float32')
+ y = net(x)
+ assert np.dtype(y.dtype).name == 'float32'
+ net.cast('float16')
+ x = x.astype('float16')
+ y1 = net(x)
+ assert np.dtype(y1.dtype).name == 'float16'
 
 if __name__ == '__main__':
  import nose