From dfe59a5f25892aa16f9e7cf84e609836144e5cbb Mon Sep 17 00:00:00 2001
From: Evan Smothers <ebs@meta.com>
Date: Wed, 18 Oct 2023 14:23:55 -0700
Subject: [PATCH] Utilities for deterministic unit tests (#497)

Summary:
Pull Request resolved: https://github.com/facebookresearch/multimodal/pull/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 | 46 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/tests/test_utils.py b/tests/test_utils.py
index abc790d95..51f94ade0 100644
--- a/tests/test_utils.py
+++ b/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,
+            )
+        )