ssk015
/
mindspore2022
forked from huawei/mindspore2022
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

clear codecheck for thor optimizer

This commit is contained in:
sl_wang 2021-05-29 19:20:18 +08:00
parent 329546b83d
commit f40fb4e157
4 changed files with 180 additions and 251 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

181
mindspore/nn/layer/thor_layer.py
View File

@ -31,21 +31,22 @@ __all__ = ['DenseThor', 'Conv2dThor', 'EmbeddingThor']
class DenseThor(Cell):
r"""
The dense connected layer.
The dense connected layer and saving the information needed for THOR.
Applies dense connected layer for the input. This layer implements the operation as:
Applies dense connected layer for the input and saves the information A and G in the dense connected layer
needed for THOR, the detail can be seen in paper: https://www.aaai.org/AAAI21Papers/AAAI-6611.ChenM.pdf
This layer implements the operation as:
.. math::
\text{outputs} = \text{activation}(\text{inputs} * \text{kernel} + \text{bias}),
where :math:`\text{activation}` is the activation function passed as the activation
argument (if passed in), :math:`\text{kernel}` is a weight matrix with the same
where :math:`\text{activation}` is the activation function , :math:`\text{kernel}` is a weight matrix with the same
data type as the inputs created by the layer, and :math:`\text{bias}` is a bias vector
with the same data type as the inputs created by the layer (only if has_bias is True).
Args:
in_channels (int): The number of channels in the input space.
out_channels (int): The number of channels in the output space.
in_channels (int): The number of the input channels.
out_channels (int): The number of the output channels.
weight_init (Union[Tensor, str, Initializer, numbers.Number]): The trainable weight_init parameter. The dtype
is same as input x. The values of str refer to the function `initializer`. Default: 'normal'.
bias_init (Union[Tensor, str, Initializer, numbers.Number]): The trainable bias_init parameter. The dtype is
@ -55,7 +56,7 @@ class DenseThor(Cell):
Default: None.
Raises:
ValueError: If weight_init or bias_init shape is incorrect.
ValueError: If weight_init shape or bias_init shape is incorrect.
Inputs:
- **input** (Tensor) - Tensor of shape :math:`(N, in\_channels)`.
@ -65,7 +66,7 @@ class DenseThor(Cell):
Examples:
>>> input = Tensor(np.random.randint(0, 255, [2, 3]), mindspore.float32)
>>> net = nn.Dense(3, 4)
>>> net = nn.DenseThor(3, 4)
>>> net(input)
[[ 2.5246444 2.2738023 0.5711005 -3.9399147 ]
[ 1.0739875 4.0155234 0.94188046 -5.459526 ]]
@ -193,40 +194,29 @@ class DenseThor(Cell):
return s
class _Conv(Cell):
class _ConvThor(Cell):
"""
Applies a N-D convolution over an input signal composed of several input planes.
Applies a N-D convolution over an input signal composed of multiple input planes.
"""
def __init__(self,
in_channels,
out_channels,
kernel_size,
stride,
pad_mode,
padding,
dilation,
group,
has_bias,
weight_init,
bias_init,
transposed=False):
super(_Conv, self).__init__()
def __init__(self, in_channels, out_channels, kernel_size, stride, pad_mode,
padding, dilation, group, has_bias, weight_init, bias_init, transposed=False):
super(_ConvThor, self).__init__()
self.in_channels = Validator.check_positive_int(in_channels)
self.out_channels = Validator.check_positive_int(out_channels)
self.kernel_size = kernel_size
self.stride = stride
self.pad_mode = pad_mode
self.bias_init = bias_init
if isinstance(padding, int):
Validator.check_non_negative_int(padding, 'padding', self.cls_name)
self.padding = padding
elif isinstance(padding, tuple):
if isinstance(padding, tuple):
for pad in padding:
Validator.check_non_negative_int(pad, 'padding item', self.cls_name)
self.padding = padding
elif isinstance(padding, int):
Validator.check_non_negative_int(padding, 'padding', self.cls_name)
self.padding = padding
else:
raise TypeError("padding type must be int/tuple(int) cannot be {}!".format(type(padding)))
raise TypeError("padding type must be int or tuple(int) cannot be {}!".format(type(padding)))
self.dilation = dilation
self.group = Validator.check_positive_int(group)
@ -235,15 +225,15 @@ class _Conv(Cell):
self._validate_stride(stride)
self._validate_dilation(dilation)
if in_channels % group != 0:
raise ValueError("Attr 'in_channels' of 'Conv2D' Op must be divisible by "
"attr 'group' of 'Conv2D' Op.")
raise ValueError("Attr 'in_channels' of 'Conv2DThor' Op must be divisible by "
"attr 'group' of 'Conv2DThor' Op.")
if out_channels % group != 0:
raise ValueError("Attr 'out_channels' of 'Conv2D' Op must be divisible by "
"attr 'group' of 'Conv2D' Op.")
if transposed:
shape = [in_channels, out_channels // group, *kernel_size]
else:
raise ValueError("Attr 'out_channels' of 'Conv2DThor' Op must be divisible by "
"attr 'group' of 'Conv2DThor' Op.")
if not transposed:
shape = [out_channels, in_channels // group, *kernel_size]
else:
shape = [in_channels, out_channels // group, *kernel_size]
self.weight = Parameter(initializer(weight_init, shape), name='weight')
if Validator.check_bool(has_bias):
@ -275,19 +265,20 @@ class _Conv(Cell):
raise ValueError("Attr 'dilation' of 'Conv2D' Op passed "
+ str(self.dilation) + ", should be a int or tuple and equal to or greater than 1.")
def construct(self, *inputs):
"""Must be overridden by all subclasses."""
raise NotImplementedError
class Conv2dThor(_Conv):
class Conv2dThor(_ConvThor):
r"""
2D convolution layer.
2D convolution layer and saving the information needed for THOR.
Applies a 2D convolution over an input tensor which is typically of shape :math:`(N, C_{in}, H_{in}, W_{in})`,
where :math:`N` is batch size, :math:`C_{in}` is channel number, and :math:`H_{in}, W_{in})` are height and width.
And saves the information A and G in the 2D convolution layer needed for THOR.
The detail can be seen in paper: https://www.aaai.org/AAAI21Papers/AAAI-6611.ChenM.pdf
For each batch of shape :math:`(C_{in}, H_{in}, W_{in})`, the formula is defined as:
.. math::
out_j = \sum_{i=0}^{C_{in} - 1} ccor(W_{ij}, X_i) + b_j,
@ -306,55 +297,51 @@ class Conv2dThor(_Conv):
:math:`\left \lfloor{1 + \frac{W_{in} + 2 \times \text{padding} - \text{ks_w} -
(\text{ks_w} - 1) \times (\text{dilation} - 1) }{\text{stride}}} \right \rfloor` respectively.
The first introduction can be found in paper `Gradient Based Learning Applied to Document Recognition
<http://vision.stanford.edu/cs598_spring07/papers/Lecun98.pdf>`_.
Args:
in_channels (int): The number of input channel :math:`C_{in}`.
out_channels (int): The number of output channel :math:`C_{out}`.
in_channels (int): The number of the input channel :math:`C_{in}`.
out_channels (int): The number of the output channel :math:`C_{out}`.
kernel_size (Union[int, tuple[int]]): The data type is int or a tuple of 2 integers. Specifies the height
and width of the 2D convolution window. Single int means the value is for both the height and the width of
the kernel. A tuple of 2 ints means the first value is for the height and the other is for the
width of the kernel.
stride (Union[int, tuple[int]]): The distance of kernel moving, an int number that represents
the height and width of movement are both strides, or a tuple of two int numbers that
represent height and width of movement respectively. Default: 1.
and width of the 2D convolution window. Single int means that the value is not only the height, but also
the width of the kernel. A tuple of 2 integers means the height and the width of the kernel respectively.
stride (Union[int, tuple[int]]): The distance of kernel moving, an int number represents the height and width
of movement, or a tuple of two int numbers that represent height and width of movement, respectively.
Default: 1.
pad_mode (str): Specifies padding mode. The optional values are
"same", "valid", "pad". Default: "same".
- same: Adopts the way of completion. The height and width of the output will be the same as
- same: Adopts the way of completion. The shape of the output will be the same as
the input. The total number of padding will be calculated in horizontal and vertical
directions and evenly distributed to top and bottom, left and right if possible. Otherwise, the
last extra padding will be done from the bottom and the right side. If this mode is set, `padding`
must be 0.
- valid: Adopts the way of discarding. The possible largest height and width of output will be returned
without padding. Extra pixels will be discarded. If this mode is set, `padding`
must be 0.
without padding. Extra pixels will be discarded. If this mode is set, `padding` must be 0.
- pad: Implicit paddings on both sides of the input. The number of `padding` will be padded to the input
Tensor borders. `padding` must be greater than or equal to 0.
padding (Union[int, tuple[int]]): Implicit paddings on both sides of the input. If `padding` is one integer,
padding (Union[int, tuple[int]]): Implicit paddings on both sides of the input. If `padding` is an integer,
the paddings of top, bottom, left and right are the same, equal to padding. If `padding` is a tuple
with four integers, the paddings of top, bottom, left and right will be equal to padding[0],
padding[1], padding[2], and padding[3] accordingly. Default: 0.
dilation (Union[int, tuple[int]]): The data type is int or a tuple of 2 integers. Specifies the dilation rate
to use for dilated convolution. If set to be :math:`k > 1`, there will
be :math:`k - 1` pixels skipped for each sampling location. Its value must
be greater or equal to 1 and bounded by the height and width of the
input. Default: 1.
be greater or equal to 1 and bounded by the height and width of the input.
Default: 1.
group (int): Splits filter into groups, `in_ channels` and `out_channels` must be
divisible by the number of groups. If the group is equal to `in_channels` and `out_channels`,
this 2D convolution layer also can be called 2D depthwise convolution layer. Default: 1.
has_bias (bool): Specifies whether the layer uses a bias vector. Default: False.
weight_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the convolution kernel.
weight_init (Union[Tensor, str, Initializer, numbers.Number]): Initializes the convolution kernel.
It can be a Tensor, a string, an Initializer or a number. When a string is specified,
values from 'TruncatedNormal', 'Normal', 'Uniform', 'HeUniform' and 'XavierUniform' distributions as well
as constant 'One' and 'Zero' distributions are possible. Alias 'xavier_uniform', 'he_uniform', 'ones'
and 'zeros' are acceptable. Uppercase and lowercase are both acceptable. Refer to the values of
Initializer for more details. Default: 'normal'.
bias_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the bias vector. Possible
bias_init (Union[Tensor, str, Initializer, numbers.Number]): Initializes the bias vector. Possible
Initializer and string are the same as 'weight_init'. Refer to the values of
Initializer for more details. Default: 'zeros'.
@ -365,48 +352,24 @@ class Conv2dThor(_Conv):
Tensor of shape :math:`(N, C_{out}, H_{out}, W_{out})`.
Examples:
>>> net = nn.Conv2d(120, 240, 4, has_bias=False, weight_init='normal')
>>> net = nn.Conv2dThor(120, 240, 4, has_bias=False, weight_init='normal')
>>> input = Tensor(np.ones([1, 120, 1024, 640]), mindspore.float32)
>>> net(input).shape
>>> print(net(input).shape)
(1, 240, 1024, 640)
"""
def __init__(self,
in_channels,
out_channels,
kernel_size,
stride=1,
pad_mode='same',
padding=0,
dilation=1,
group=1,
has_bias=False,
weight_init='normal',
bias_init='zeros'):
def __init__(self, in_channels, out_channels, kernel_size, stride=1,
pad_mode='same', padding=0, dilation=1, group=1, has_bias=False,
weight_init='normal', bias_init='zeros'):
kernel_size = twice(kernel_size)
stride = twice(stride)
self._dilation = dilation
dilation = twice(dilation)
super(Conv2dThor, self).__init__(
in_channels,
out_channels,
kernel_size,
stride,
pad_mode,
padding,
dilation,
group,
has_bias,
weight_init,
bias_init)
self.conv2d = P.Conv2D(out_channel=self.out_channels,
kernel_size=self.kernel_size,
mode=1,
pad_mode=self.pad_mode,
pad=self.padding,
stride=self.stride,
dilation=self.dilation,
group=self.group)
super(Conv2dThor, self).__init__(in_channels, out_channels, kernel_size,
stride, pad_mode, padding, dilation, group, has_bias, weight_init, bias_init)
self.conv2d = P.Conv2D(out_channel=self.out_channels, kernel_size=self.kernel_size,
mode=1, pad_mode=self.pad_mode, pad=self.padding,
stride=self.stride, dilation=self.dilation, group=self.group)
self._init_depthwise_conv2d(weight_init)
self.bias_add = P.BiasAdd()
@ -550,40 +513,32 @@ class Conv2dThor(_Conv):
return output
def extend_repr(self):
s = 'input_channels={}, output_channels={}, kernel_size={},' \
'stride={}, pad_mode={}, padding={}, dilation={}, ' \
'group={}, has_bias={},' \
'weight_init={}, bias_init={}'.format(
self.in_channels,
self.out_channels,
self.kernel_size,
self.stride,
self.pad_mode,
self.padding,
self.dilation,
self.group,
self.has_bias,
self.weight_init,
self.bias_init)
s = 'input_channels={}, output_channels={}, kernel_size={},' 'stride={}, ' \
'pad_mode={}, padding={}, dilation={}, ' 'group={}, has_bias={},' \
'weight_init={}, bias_init={}'.format(self.in_channels, self.out_channels, self.kernel_size,
self.stride, self.pad_mode, self.padding, self.dilation,
self.group, self.has_bias, self.weight_init, self.bias_init)
return s
class EmbeddingThor(Cell):
r"""
A simple lookup table that stores embeddings of a fixed dictionary and size.
A simple lookup table that stores embeddings of a fixed dictionary and size
and saving the information needed for THOR.
This module is often used to store word embeddings and retrieve them using
indices. The input to the module is a list of indices, and the output is
the corresponding word embeddings.
the corresponding word embeddings. And saves the information A and G in the dense connected layer
needed for THOR, the detail can be seen in paper: https://www.aaai.org/AAAI21Papers/AAAI-6611.ChenM.pdf
Note:
When 'use_one_hot' is set to True, the type of the input must be mindspore.int32.
Args:
vocab_size (int): Size of the dictionary of embeddings.
vocab_size (int): The size of the dictionary of embeddings.
embedding_size (int): The size of each embedding vector.
use_one_hot (bool): Specifies whether to apply one_hot encoding form. Default: False.
embedding_table (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the embedding_table.
embedding_table (Union[Tensor, str, Initializer, numbers.Number]): Initializes the embedding_table.
Refer to class `initializer` for the values of string when a string
is specified. Default: 'normal'.
dtype (:class:`mindspore.dtype`): Data type of input. Default: mindspore.float32.

196
mindspore/nn/optim/thor.py
View File

@ -20,7 +20,6 @@ from mindspore.common.parameter import Parameter, ParameterTuple
from mindspore.common.tensor import Tensor
import mindspore.nn as nn
import mindspore.common.dtype as mstype
from mindspore.nn.cell import Cell
from mindspore._checkparam import Validator
from mindspore.nn.optim.optimizer import Optimizer
from mindspore.parallel._utils import _get_device_num, _get_gradients_mean
@ -87,56 +86,12 @@ def _clip_grad(clip_type, clip_value, grad):
new_grad = nn.ClipByNorm()(grad, F.cast(F.tuple_to_array((clip_value,)), dt))
return new_grad
get_square_sum = C.MultitypeFuncGraph("get_square_sum")
@get_square_sum.register("Tensor")
def _get_square_sum(grad):
norm = P.ReduceSum(False)(F.square(grad), ())
norm = F.expand_dims(F.cast(norm, mstype.float32), 0)
return norm
apply_global_norm = C.MultitypeFuncGraph("apply_global_norm")
@apply_global_norm.register("Tensor", "Tensor", "Tensor")
def _apply_global_norm(clip_norm, global_norm, grad):
grad = grad * clip_norm / global_norm
return grad
class GlobalNorm(Cell):
"""
Calculate the global norm value of given tensors
"""
def __init__(self):
super(GlobalNorm, self).__init__()
self.norm = nn.Norm()
self.hyper_map = C.HyperMap()
def construct(self, grads):
square_sum = self.hyper_map(get_square_sum, grads)
global_norms = F.sqrt(F.addn(square_sum) / F.scalar_to_array(len(square_sum)))
return global_norms
class ClipByGlobalNorm(Cell):
"""
Clip grads by global norm
"""
def __init__(self, clip_norm=1.0):
super(ClipByGlobalNorm, self).__init__()
self.global_norm = GlobalNorm()
self.clip_norm = Tensor([clip_norm], mstype.float32)
self.hyper_map = C.HyperMap()
def construct(self, grads):
global_norm = self.global_norm(grads)
cond = P.GreaterEqual()(global_norm, self.clip_norm)
global_norm = F.select(cond, global_norm, self.clip_norm)
grads = self.hyper_map(F.partial(apply_global_norm, self.clip_norm, global_norm), grads)
return grads
def clip_gradient(enable_clip_grad, gradients):
"""clip gradients"""
if enable_clip_grad:
if IS_ENABLE_GLOBAL_NORM:
gradients = ClipByGlobalNorm()(gradients)
gradients = C.clip_by_global_norm(gradients, GRADIENT_CLIP_VALUE, None)
else:
gradients = hyper_map_op(F.partial(clip_grad, GRADIENT_CLIP_TYPE, GRADIENT_CLIP_VALUE), gradients)
return gradients
@ -211,16 +166,15 @@ def get_net_layertype_mask(net):
def get_layer_counter(layer_type, layer_counter, params, idx):
"""get layer counter"""
if layer_type in [Conv, FC, LayerNorm, BatchNorm]:
if layer_type in [LayerNorm, BatchNorm]:
if "beta" in params[idx].name.lower():
layer_counter = layer_counter + 1
if layer_type in [Conv, FC]:
if "bias" in params[idx].name.lower():
layer_counter = layer_counter + 1
else:
if "bias" in params[idx].name.lower():
if idx < len(params) - 1 and "bias" not in params[idx + 1].name.lower():
layer_counter = layer_counter + 1
else:
if idx < len(params) - 1 and "bias" not in params[idx + 1].name.lower():
layer_counter = layer_counter + 1
elif layer_type in [LayerNorm, BatchNorm]:
if "beta" in params[idx].name.lower():
layer_counter = layer_counter + 1
else:
layer_counter = layer_counter + 1
return layer_counter
@ -229,6 +183,12 @@ def get_layer_counter(layer_type, layer_counter, params, idx):
def thor(net, learning_rate, damping, momentum, weight_decay=0.0, loss_scale=1.0, batch_size=32,
use_nesterov=False, decay_filter=lambda x: x.name not in [], split_indices=None, enable_clip_grad=False,
frequency=100):
r"""
Updates gradients by the THOR algorithm.
Trace-based Hardware-driven layer-ORiented Natural Gradient Descent Computation (THOR) algorithm is proposed in:
`THOR: Trace-based Hardware-driven layer-ORiented Natural Gradient Descent Computation
<https://www.aaai.org/AAAI21Papers/AAAI-6611.ChenM.pdf>`_
"""
context.set_context(max_call_depth=10000)
ConvertNetUntils().convert_to_thor_net(net)
if context.get_context("device_target") == "Ascend":
@ -285,8 +245,6 @@ class ThorGpu(Optimizer):
self.update_gradient = P.UpdateThorGradient(split_dim=split_dim)
self.enable_clip_grad = enable_clip_grad
self.frequency = frequency
self.parallel_mode = context.get_auto_parallel_context("parallel_mode")
self.is_distributed = (self.parallel_mode != ParallelMode.STAND_ALONE)
self._define_gpu_reducer(split_indices)
def get_frequency(self):
@ -320,6 +278,8 @@ class ThorGpu(Optimizer):
def _define_gpu_reducer(self, split_indices):
"""define gpu reducer"""
self.parallel_mode = context.get_auto_parallel_context("parallel_mode")
self.is_distributed = (self.parallel_mode != ParallelMode.STAND_ALONE)
if self.is_distributed:
mean = _get_gradients_mean()
degree = _get_device_num()
@ -369,21 +329,21 @@ class ThorGpu(Optimizer):
if layer_type in [Conv, FC, Embedding] and "bias" not in self.params[idx].name.lower():
self.weight_fim_idx_map = self.weight_fim_idx_map + (self.thor_layer_count,)
self.weight_layertype_idx_map = self.weight_layertype_idx_map + (layer_type,)
self.thor_layer_count = self.thor_layer_count + 1
self.weight_layertype_idx_map = self.weight_layertype_idx_map + (layer_type,)
if layer_type == Conv:
self.weight_conv_idx_map = self.weight_conv_idx_map + (self.conv_layer_count,)
self.conv_layer_count = self.conv_layer_count + 1
else:
self.weight_conv_idx_map = self.weight_conv_idx_map + (-1,)
else:
self.weight_fim_idx_map = self.weight_fim_idx_map + (-1,)
self.weight_conv_idx_map = self.weight_conv_idx_map + (-1,)
self.weight_fim_idx_map = self.weight_fim_idx_map + (-1,)
if layer_type == LayerNorm:
self.weight_layertype_idx_map = self.weight_layertype_idx_map + (LayerNorm,)
else:
self.weight_layertype_idx_map = self.weight_layertype_idx_map + (Other,)
# bert.cls1.output_bias: not a network layer, only a trainable param
# bert.cls1.output_bias: not a network layer, only a trainable param
if "output_bias" not in self.params[idx].name.lower():
layer_counter = get_layer_counter(layer_type, layer_counter, self.params, idx)
@ -439,6 +399,24 @@ class ThorGpu(Optimizer):
matrix_g_allreduce = matrix_g_allreduce + (matrix_g,)
return matrix_a_allreduce, matrix_g_allreduce
def _process_layernorm(self, damping_step, gradient):
"""process layernorm"""
damping = self.sqrt(damping_step)
normalizer = self.batch_size
normalizer = self.cast(normalizer, mstype.float32)
fim_cov = self.square(gradient)
fim_cov = self.mul(fim_cov, 1.0 / normalizer)
fim_cov = fim_cov + damping
fim_inv = self.inv(fim_cov)
gradient = self.mul(fim_inv, gradient)
return gradient
def _reshape_gradient(self, conv_layer_count, g, g_shape):
"""reshape gradient"""
if conv_layer_count != -1:
g = self.reshape(g, g_shape)
return g
def construct(self, gradients):
params = self.params
moments = self.moments
@ -470,8 +448,7 @@ class ThorGpu(Optimizer):
fake_g = self.assign(self.matrix_g[thor_layer_count], matrix_g)
g = F.depend(g, fake_a)
g = F.depend(g, fake_g)
if conv_layer_count != -1:
g = self.reshape(g, g_shape)
g = self._reshape_gradient(conv_layer_count, g, g_shape)
elif layer_type == Embedding:
matrix_a = matrix_a_allreduce[thor_layer_count]
matrix_g = matrix_g_allreduce[thor_layer_count]
@ -483,14 +460,7 @@ class ThorGpu(Optimizer):
g = self.mul(temp_a, g)
g = self.matmul(g, matrix_g)
elif layer_type == LayerNorm:
damping = self.sqrt(damping_step)
normalizer = self.batch_size
normalizer = self.cast(normalizer, mstype.float32)
fim_cov = self.square(g)
fim_cov = self.mul(fim_cov, 1.0 / normalizer)
fim_cov = fim_cov + damping
fim_inv = self.inv(fim_cov)
g = self.mul(fim_inv, g)
g = self._process_layernorm(damping_step, g)
new_grads = new_grads + (g,)
else:
for j in range(len(self.params)):
@ -504,8 +474,7 @@ class ThorGpu(Optimizer):
matrix_a = self.matrix_a[thor_layer_count]
matrix_g = self.matrix_g[thor_layer_count]
g = self.update_gradient(matrix_g, g, matrix_a)
if conv_layer_count != -1:
g = self.reshape(g, g_shape)
g = self._reshape_gradient(conv_layer_count, g, g_shape)
elif layer_type == Embedding:
matrix_a = self.matrix_a[thor_layer_count]
matrix_g = self.matrix_g[thor_layer_count]
@ -514,14 +483,7 @@ class ThorGpu(Optimizer):
g = self.mul(temp_a, g)
g = self.matmul(g, matrix_g)
elif layer_type == LayerNorm:
damping = self.sqrt(damping_step)
normalizer = self.batch_size
normalizer = self.cast(normalizer, mstype.float32)
fim_cov = self.square(g)
fim_cov = self.mul(fim_cov, 1.0 / normalizer)
fim_cov = fim_cov + damping
fim_inv = self.inv(fim_cov)
g = self.mul(fim_inv, g)
g = self._process_layernorm(damping_step, g)
new_grads = new_grads + (g,)
gradients = new_grads
@ -564,8 +526,8 @@ class ThorAscend(Optimizer):
self.matrix_g = ()
self.thor_layer_count = 0
self.conv_layer_count = 0
self.weight_fim_idx_map = ()
self.weight_conv_idx_map = ()
self.weight_fim_idx_map = ()
self.weight_layertype_idx_map = ()
self._process_matrix_init_and_weight_idx_map(self.net)
self.matrix_a = ParameterTuple(self.matrix_a)
@ -584,8 +546,6 @@ class ThorAscend(Optimizer):
self.batch_size_scale = Tensor(batch_size * batch_size, mstype.float32)
self.enable_clip_grad = enable_clip_grad
self.frequency = frequency
self.parallel_mode = context.get_auto_parallel_context("parallel_mode")
self.is_distributed = (self.parallel_mode != ParallelMode.STAND_ALONE)
self._define_ascend_reducer(split_indices)
@ -626,6 +586,8 @@ class ThorAscend(Optimizer):
def _define_ascend_reducer(self, split_indices):
"""define ascend reducer"""
self.parallel_mode = context.get_auto_parallel_context("parallel_mode")
self.is_distributed = (self.parallel_mode != ParallelMode.STAND_ALONE)
if self.is_distributed:
mean = _get_gradients_mean()
degree = _get_device_num()
@ -937,6 +899,44 @@ class ThorAscend(Optimizer):
new_grads = new_grads + (g,)
return new_grads
def _get_second_grad_by_matmul(self, index, temp_a, temp_g, g, temp_max):
"""get second gradient by matmul"""
conv_layer_count = self.weight_conv_idx_map[index]
layer_type = self.weight_layertype_idx_map[index]
if layer_type == FC:
g = self.cube_matmul_left_fc(temp_g, g)
g = self.cube_matmul_right_fc(g, temp_a, temp_max)
elif layer_type == Conv:
a_normalizer = self.a_normalizer[conv_layer_count]
a_normalizer = F.depend(a_normalizer, g)
temp_max = self.mul(temp_max, self.batch_size / a_normalizer)
g = self.cube_matmul_left(temp_g, g)
g = self.cube_matmul_right_mul(g, temp_a, temp_max)
return g, temp_max
def _get_second_grad_by_layertype(self, index, matrix_a_allreduce, matrix_g_allreduce, g, damping_step):
"""get second gradient by layertype"""
thor_layer_count = self.weight_fim_idx_map[index]
layer_type = self.weight_layertype_idx_map[index]
if layer_type == Embedding:
temp_a_ori = matrix_a_allreduce[thor_layer_count]
temp_g = matrix_g_allreduce[thor_layer_count]
fake_a = self.assign(self.matrix_a_cov[thor_layer_count], temp_a_ori)
fake_g = self.assign(self.matrix_g_cov[thor_layer_count], temp_g)
g = F.depend(g, fake_a)
g = F.depend(g, fake_g)
temp_a = self.expand(temp_a_ori, 1)
g = self.mul(temp_a, g)
temp_g = self.cast(temp_g, mstype.float16)
g = self.cast(g, mstype.float16)
g = self.matmul(g, temp_g)
g = self.cast(g, mstype.float32)
elif layer_type == FC:
g = self._process_thor_fc(thor_layer_count, matrix_a_allreduce, matrix_g_allreduce, g)
elif layer_type == LayerNorm:
g = self._process_layernorm(damping_step, g)
return g
def construct(self, gradients):
params = self.params
moments = self.moments
@ -963,8 +963,6 @@ class ThorAscend(Optimizer):
for i in range(len(self.params)):
g = gradients[i]
thor_layer_count = self.weight_fim_idx_map[i]
conv_layer_count = self.weight_conv_idx_map[i]
layer_type = self.weight_layertype_idx_map[i]
temp_a = matrix_a_allreduce[thor_layer_count]
temp_g = matrix_g_allreduce[thor_layer_count]
matrix_a_inv_max = self.log(matrix_a_max_allreduce[thor_layer_count])
@ -979,15 +977,7 @@ class ThorAscend(Optimizer):
matrix_g_max_allreduce[thor_layer_count])
temp_a = self.cast(temp_a, mstype.float16)
temp_g = self.cast(temp_g, mstype.float16)
if layer_type == FC:
g = self.cube_matmul_left_fc(temp_g, g)
g = self.cube_matmul_right_fc(g, temp_a, temp_max)
elif layer_type == Conv:
a_normalizer = self.a_normalizer[conv_layer_count]
a_normalizer = F.depend(a_normalizer, g)
temp_max = self.mul(temp_max, self.batch_size / a_normalizer)
g = self.cube_matmul_left(temp_g, g)
g = self.cube_matmul_right_mul(g, temp_a, temp_max)
g, temp_max = self._get_second_grad_by_matmul(i, temp_a, temp_g, g, temp_max)
fake_a = self.assign(self.matrix_a[thor_layer_count], temp_a)
fake_g = self.assign(self.matrix_g[thor_layer_count], temp_g)
fake_max = self.assign(self.matrix_max_inv[thor_layer_count], temp_max)
@ -999,25 +989,7 @@ class ThorAscend(Optimizer):
else:
for i in range(len(self.params)):
g = gradients[i]
thor_layer_count = self.weight_fim_idx_map[i]
layer_type = self.weight_layertype_idx_map[i]
if layer_type == Embedding:
temp_a_ori = matrix_a_allreduce[thor_layer_count]
temp_g = matrix_g_allreduce[thor_layer_count]
fake_a = self.assign(self.matrix_a_cov[thor_layer_count], temp_a_ori)
fake_g = self.assign(self.matrix_g_cov[thor_layer_count], temp_g)
g = F.depend(g, fake_a)
g = F.depend(g, fake_g)
temp_a = self.expand(temp_a_ori, 1)
g = self.mul(temp_a, g)
temp_g = self.cast(temp_g, mstype.float16)
g = self.cast(g, mstype.float16)
g = self.matmul(g, temp_g)
g = self.cast(g, mstype.float32)
elif layer_type == FC:
g = self._process_thor_fc(thor_layer_count, matrix_a_allreduce, matrix_g_allreduce, g)
elif layer_type == LayerNorm:
g = self._process_layernorm(damping_step, g)
g = self._get_second_grad_by_layertype(i, matrix_a_allreduce, matrix_g_allreduce, g, damping_step)
new_grads = new_grads + (g,)
gradients = new_grads
else:

44
mindspore/ops/_op_impl/_custom_op/batch_matmul_impl.py
View File

@ -43,8 +43,9 @@ def _get_flattern_shape(shape):
return (flattern_shape,)
def _inner_matmul_new(tik_instance, dtype, input1, input1_index, input2, input2_index, res, res_index):
def _inner_matmul_new(tik_instance, dtype, input_info, res, res_index):
"""_inner_matmul_new"""
input1, input1_index, input2, input2_index = input_info
input_1_local_ub = tik_instance.Tensor(dtype, [128], name="input_1_local_ub", scope=tik.scope_ubuf)
t_1_0_local_ub = tik_instance.Tensor(dtype, [64 * 128], name="t_1_0_local_ub", scope=tik.scope_ubuf)
tik_instance.data_move(input_1_local_ub, input1[input1_index], 0, 1, 16, 0, 0)
@ -75,8 +76,9 @@ def _inner_matmul_new(tik_instance, dtype, input1, input1_index, input2, input2_
matmul_hybrid_f_t_local_ub, 0, 1, 8, 0, 0)
def _inner_matmul_new_1_64_32_64(tik_instance, dtype, input1, input1_index, input2, input2_index, res, res_index):
def _inner_matmul_new_1_64_32_64(tik_instance, dtype, input_info, res, res_index):
"""_inner_matmul_new_1_64_32_64"""
input1, input1_index, input2, input2_index = input_info
input_1_local_ub = tik_instance.Tensor(dtype, [64], name="input_1_local_ub", scope=tik.scope_ubuf)
tik_instance.data_move(input_1_local_ub, input1[input1_index], 0, 1, 8, 0, 0)
with tik_instance.for_range(0, 2, thread_num=2) as thread_idx2:
@ -94,8 +96,9 @@ def _inner_matmul_new_1_64_32_64(tik_instance, dtype, input1, input1_index, inpu
matmul_hybrid_f_t_local_ub, 0, 1, 4, 0, 0)
def process_input_shape_640(input_shape, tik_instance, dtype, input1, input2, res):
def process_input_shape_640(input_shape, tik_instance, dtype, total_input, res):
"""process input shape of 640"""
input1, input2 = total_input
if input_shape == ((5, 128, 128), (5, 128, 128), "float32", False, True):
with tik_instance.for_range(0, 30, block_num=30) as block_idx,\
tik_instance.for_range(0, 11) as cc1_db,\
@ -144,25 +147,26 @@ def process_input_shape_640(input_shape, tik_instance, dtype, input1, input2, re
matmul_hybrid_f_t_local_ub, 0, 1, 8, 0, 0)
def process_input_shape_1152(input_shape, tik_instance, dtype, input1, input2, res):
def process_input_shape_1152(input_shape, tik_instance, dtype, total_input, res):
"""process input shape of 1152"""
input1, input2 = total_input
if input_shape == ((9, 128, 128), (9, 128, 128), "float32", False, True):
with tik_instance.for_range(0, 27, block_num=27) as block_idx:
with tik_instance.for_range(0, 42, thread_num=2) as cc0:
input1_index = (block_idx // 3) * 16384 + (block_idx % 3) * 5504 + cc0 * 128
input2_index = (block_idx // 3) * 16384
res_index = (block_idx // 3) * 16384 + (block_idx % 3) * 5504 + cc0 * 128
input_info = input1, input1_index, input2, input2_index
_inner_matmul_new(tik_instance, dtype,
input1, input1_index,
input2, input2_index,
input_info,
res, res_index)
with tik_instance.if_scope((block_idx % 3) < 2):
input1_index = (block_idx // 3) * 16384 + (block_idx % 3) * 5504 + 42 * 128
input2_index = (block_idx // 3) * 16384
res_index = (block_idx // 3) * 16384 + (block_idx % 3) * 5504 + 42 * 128
input_info = input1, input1_index, input2, input2_index
_inner_matmul_new(tik_instance, dtype,
input1, input1_index,
input2, input2_index,
input_info,
res, res_index)
@ -216,12 +220,12 @@ def cus_batch_matmul(input_x1, input_x2, output, transpose_a=False,
input1_index = block_idx * 32768 + cc0 * 16384 + cc1 * 128
input2_index = block_idx * 32768 + cc0 * 16384
res_index = block_idx * 32768 + cc0 * 16384 + cc1 * 128
input_info = input1, input1_index, input2, input2_index
_inner_matmul_new(tik_instance, dtype,
input1, input1_index,
input2, input2_index,
res, res_index)
input_info, res, res_index)
process_input_shape_640(input_shape, tik_instance, dtype, input1, input2, res)
total_input = input1, input2
process_input_shape_640(input_shape, tik_instance, dtype, total_input, res)
if input_shape == ((18, 128, 128), (18, 128, 128), "float32", False, True):
with tik_instance.for_range(0, 18, block_num=18) as block_idx, \
@ -229,12 +233,11 @@ def cus_batch_matmul(input_x1, input_x2, output, transpose_a=False,
input1_index = block_idx * 16384 + cc0 * 128
input2_index = block_idx * 16384
res_index = block_idx * 16384 + cc0 * 128
input_info = input1, input1_index, input2, input2_index
_inner_matmul_new(tik_instance, dtype,
input1, input1_index,
input2, input2_index,
res, res_index)
input_info, res, res_index)
process_input_shape_1152(input_shape, tik_instance, dtype, input1, input2, res)
process_input_shape_1152(input_shape, tik_instance, dtype, total_input, res)
if input_shape == ((1, 64, 64), (1, 64, 64), "float32", False, True):
with tik_instance.for_range(0, 32, block_num=32) as block_idx,\
@ -242,9 +245,9 @@ def cus_batch_matmul(input_x1, input_x2, output, transpose_a=False,
input1_index = block_idx * 128 + cc0 * 64
input2_index = 0
res_index = block_idx * 128 + cc0 * 64
input_info = input1, input1_index, input2, input2_index
_inner_matmul_new_1_64_32_64(tik_instance, dtype,
input1, input1_index,
input2, input2_index,
input_info,
res, res_index)
input_shape_list = [((1, 128, 128), (1, 128, 128), "float32", False, True),
@ -271,9 +274,8 @@ def cus_batch_matmul(input_x1, input_x2, output, transpose_a=False,
else:
input2_index = 0
res_index = block_idx * block_process_ele_num + cc0 * input1_unit_size
_inner_matmul_new(tik_instance, dtype,
input1, input1_index,
input2, input2_index,
input_info = input1, input1_index, input2, input2_index
_inner_matmul_new(tik_instance, dtype, input_info,
res, res_index)
tik_instance.BuildCCE(kernel_name, inputs=[input1, input2], outputs=[res])

10
mindspore/train/train_thor/model_thor.py
View File

@ -27,9 +27,9 @@ from mindspore._c_expression import init_exec_dataset
from .dataset_helper import DatasetHelper
def _convert_type(types):
def _convert_to_ms_type(types):
"""
Convert from numpy type to tensor type.
Convert from numpy type to mindspore tensor type.
Args:
types (list): Numpy type list of element in dataset.
@ -38,15 +38,15 @@ def _convert_type(types):
list, list of element in dataset.
"""
ms_types = []
for np_type in types:
ms_type = pytype_to_dtype(np_type)
for numpy_type in types:
ms_type = pytype_to_dtype(numpy_type)
ms_types.append(ms_type)
return ms_types
def _get_types_and_shapes(dataset):
"""Get dataset types and shapes."""
dataset_types = _convert_type(dataset.output_types())
dataset_types = _convert_to_ms_type(dataset.output_types())
dataset_shapes = dataset.output_shapes()
return dataset_types, dataset_shapes