Utilities for deterministic unit tests (facebookresearch#497)

Summary:
Pull Request resolved: facebookresearch#497

A longstanding pain point of ours has been writing unit tests that

(a) are nontrivial enough to catch any regressions, and
(b) don't break every time we add a new test case, change model initialization order, or perform other superficial refactors.

Random initialization usually passes (a) but fails (b), while deterministic initialization ([e.g.](https://github.com/facebookresearch/multimodal/blob/2ddb8cdb205f2035e88e4fafb7e88cccb7b99705/tests/test_utils.py#L192)) does the opposite.

This diff introduces a utility for constructing deterministic but nontrivial tensors of any shape, with flexibility so the user can determine the tensor's range. As an easy extension, we also add a utility to initialize nn.Module parameters in the same way.

Reviewed By: abhinavarora

Differential Revision: D50251029

fbshipit-source-id: 26a95a30cd55ac01b04dbfb687955f35e8202de5

tests/test_utils.py

@@ -11,6 +11,8 @@
 from pathlib import Path
 from typing import Any, Dict, NamedTuple, Optional, Tuple, Union
 
+import numpy as np
+
 import pytest
 import torch
 import torch.distributed as dist
@@ -240,3 +242,47 @@ def split_tensor_for_distributed_test(
  if move_to_device:
  x = x.to(device=device_id)
  return x
+
+
+def fixed_init_tensor(
+ shape: torch.Size,
+ min_val: Union[float, int] = 0.0,
+ max_val: Union[float, int] = 1.0,
+ nonlinear: bool = False,
+ dtype: torch.dtype = torch.float,
+):
+ """
+ Utility for generating deterministic tensors of a given shape. In general stuff
+ like torch.ones, torch.eye, etc can result in trivial outputs. This utility
+ generates a range tensor [min_val, max_val) of a specified dtype, applies
+ a sine function if nonlinear=True, then reshapes to the appropriate shape.
+ """
+ n_elements = np.prod(shape)
+ step_size = (max_val - min_val) / n_elements
+ x = torch.arange(min_val, max_val, step_size, dtype=dtype).reshape(shape)
+ if nonlinear:
+ return torch.sin(x)
+ return x
+
+
+@torch.no_grad
+def fixed_init_model(
+ model: nn.Module,
+ min_val: Union[float, int] = 0.0,
+ max_val: Union[float, int] = 1.0,
+ nonlinear: bool = False,
+):
+ """
+ This utility initializes all parameters of a model deterministically using the
+ function fixed_init_tensor above. See that docstring for details of each parameter.
+ """
+ for name, param in model.named_parameters():
+ param.copy_(
+ fixed_init_tensor(
+ param.shape,
+ min_val=min_val,
+ max_val=max_val,
+ nonlinear=nonlinear,
+ dtype=param.dtype,
+ )
+ )