| 217 | data_strategy=st.data(), |
| 218 | ) |
| 219 | def test_fp16_max_norm(self, row_dim, norm, data_strategy): |
| 220 | weight = np.random.rand(row_dim, 5).astype(np.float16) |
| 221 | grad = np.random.rand(row_dim, 5).astype(np.float16) |
| 222 | |
| 223 | # generate indices that will be updated |
| 224 | indices = data_strategy.draw( |
| 225 | hu.tensor( |
| 226 | dtype=np.int64, |
| 227 | min_dim=1, |
| 228 | max_dim=1, |
| 229 | elements=st.sampled_from(np.arange(weight.shape[0])), |
| 230 | ) |
| 231 | ) |
| 232 | indices = np.unique(indices) |
| 233 | |
| 234 | # compute expected result |
| 235 | result = weight.copy() |
| 236 | # prevent dived by zero |
| 237 | eps = 1e-12 |
| 238 | norms = np.sqrt(np.sum(result[indices, ] ** 2, axis=1, keepdims=True)) |
| 239 | # if the norms are smaller than max_norm, then it doesn't need update |
| 240 | desired = np.clip(norms, 0, norm) |
| 241 | # apply max norm |
| 242 | result[indices, ] *= desired / (eps + norms) |
| 243 | |
| 244 | weight_blob = core.BlobReference("weight_blob") |
| 245 | workspace.FeedBlob(weight_blob, weight) |
| 246 | grad_blob = core.BlobReference("grad_blob") |
| 247 | workspace.FeedBlob(grad_blob, grad) |
| 248 | indices_blob = core.BlobReference("indices") |
| 249 | workspace.FeedBlob(indices_blob, indices) |
| 250 | grad_blob_slice = core.GradientSlice(indices=indices_blob, values=grad_blob) |
| 251 | train_init_net, train_net = self.get_training_nets() |
| 252 | reg = regularizer.MaxNorm(norm, dtype='fp16') |
| 253 | reg( |
| 254 | train_net, train_init_net, weight_blob, grad_blob_slice, by=RegularizationBy.AFTER_OPTIMIZER |
| 255 | ) |
| 256 | workspace.RunNetOnce(train_init_net) |
| 257 | workspace.RunNetOnce(train_net) |
| 258 | npt.assert_almost_equal(result, workspace.FetchBlob('weight_blob'), decimal=2) |