| @@ -41,7 +41,7 @@ class _BatchNorm(Cell): | |||
| moving_mean_init='zeros', | |||
| moving_var_init='ones', | |||
| use_batch_statistics=True, | |||
| group=1): | |||
| device_num_each_group=1): | |||
| super(_BatchNorm, self).__init__() | |||
| if num_features < 1: | |||
| raise ValueError("num_features must be at least 1") | |||
| @@ -60,7 +60,7 @@ class _BatchNorm(Cell): | |||
| gamma_init, num_features), name="gamma", requires_grad=affine) | |||
| self.beta = Parameter(initializer( | |||
| beta_init, num_features), name="beta", requires_grad=affine) | |||
| self.group = check_int_positive(group) | |||
| self.group = check_int_positive(device_num_each_group) | |||
| if self.group != 1: | |||
| self.rank_id = get_rank() | |||
| self.rank_size = get_group_size() | |||
| @@ -324,7 +324,7 @@ class GlobalBatchNorm(_BatchNorm): | |||
| Args: | |||
| num_features (int): `C` from an expected input of size (N, C, H, W). | |||
| group (int): The number of device in each group. | |||
| device_num_each_group (int): The number of device in each group. | |||
| eps (float): A value added to the denominator for numerical stability. Default: 1e-5. | |||
| momentum (float): A floating hyperparameter of the momentum for the | |||
| running_mean and running_var computation. Default: 0.9. | |||
| @@ -364,7 +364,7 @@ class GlobalBatchNorm(_BatchNorm): | |||
| moving_mean_init='zeros', | |||
| moving_var_init='ones', | |||
| use_batch_statistics=True, | |||
| group=1): | |||
| device_num_each_group=1): | |||
| super(GlobalBatchNorm, self).__init__(num_features, | |||
| eps, | |||
| momentum, | |||
| @@ -374,8 +374,8 @@ class GlobalBatchNorm(_BatchNorm): | |||
| moving_mean_init, | |||
| moving_var_init, | |||
| use_batch_statistics, | |||
| group) | |||
| self.group = check_int_positive(group) | |||
| device_num_each_group) | |||
| self.group = check_int_positive(device_num_each_group) | |||
| if self.group <= 1: | |||
| raise ValueError("the number of group must be greater than 1.") | |||
| def _check_data_dim(self, x): | |||
| @@ -482,17 +482,17 @@ class GroupNorm(Cell): | |||
| >>> x = Tensor(np.ones([1, 64, 256, 256], np.float32)) | |||
| >>> goup_norm_op(x) | |||
| """ | |||
| def __init__(self, num_groups, num_channels, eps=1e-05, affine=True): | |||
| def __init__(self, num_groups, num_channels, eps=1e-05, affine=True, gamma_init='ones', beta_init='zeros'): | |||
| super(GroupNorm, self).__init__() | |||
| self.num_groups = check_int_positive(num_groups) | |||
| self.num_channels = check_int_positive(num_channels) | |||
| if num_channels % num_groups != 0: | |||
| raise ValueError("num_channels should be divided by num_groups") | |||
| self.eps = Tensor(check_typename('eps', eps, (float,)), mstype.float32) | |||
| self.eps = check_typename('eps', eps, (float,)) | |||
| self.affine = check_bool(affine) | |||
| gamma = initializer('ones', [num_channels, 1, 1], mstype.float32) | |||
| beta = initializer('zeros', [num_channels, 1, 1], mstype.float32) | |||
| gamma = initializer(gamma_init, [num_channels, 1, 1]) | |||
| beta = initializer(beta_init, [num_channels, 1, 1]) | |||
| if self.affine: | |||
| self.gamma = Parameter(gamma, name='gamma') | |||
| self.beta = Parameter(beta, name='beta') | |||