Merge pull request #14 from HumanCompatibleAI/layer_nmf_types

Add type annotations to stuff

Author: dfilan

ci/code_checks.sh

@@ -6,5 +6,6 @@ SRC_FILES=(src/ tests/ setup.py)
 set -x  # echo commands
 set -e  # quit immediately on error
 
+pytype ${SRC_FILES[@]}
 flake8 ${SRC_FILES[@]}
-black --check ${SRC_FILES[@]}
+black --check ${SRC_FILES[@]}

src/reward_preprocessing/common/utils.py

@@ -8,6 +8,7 @@
 import torch as th
 from torch import nn as nn
 from torch.utils import data as torch_data
+import vegans.utils
 
 
 def make_transition_to_tensor(num_acts):
@@ -206,3 +207,9 @@ def forward(self, transition_tensor: th.Tensor) -> th.Tensor:
 
         dones = th.zeros_like(obs[:, 0])
         return self.rew_net(state=obs, action=act, next_state=next_obs, done=dones)
+
+
+def save_loss_plots(losses, save_dir):
+    """Save plots of generator/adversary losses over training."""
+    fig, _ = vegans.utils.plot_losses(losses, show=False)
+    fig.savefig(Path(save_dir) / "loss_fig.png")

src/reward_preprocessing/scripts/train_gan.py

@@ -101,7 +101,7 @@ def train_gan(
     # save samples, return losses, save plot of losses
     samples, losses = gan.get_training_results()
     utils.save_loss_plots(losses, gan.folder)
-    utils.visualize_samples(samples, num_acts, gan.folder)
+    utils.visualize_samples(samples, gan.folder)
     return losses
 
 

src/reward_preprocessing/vis/reward_vis.py

@@ -1,7 +1,7 @@
 """Port of lucid.scratch.rl_util to PyTorch. APL2.0 licensed."""
 from functools import reduce
 import logging
-from typing import List, Optional
+from typing import Callable, Dict, List, Optional, Union
 
 import lucent.optvis.param as param
 import lucent.optvis.render as render
@@ -15,7 +15,13 @@
 import reward_preprocessing.vis.objectives as objectives_rfi
 
 
-def argmax_nd(x: np.ndarray, axes: List[int], *, max_rep=np.inf, max_rep_strict=None):
+def argmax_nd(
+    x: np.ndarray,
+    axes: List[int],
+    *,
+    max_rep: Union[int, float] = np.inf,
+    max_rep_strict: Optional[bool] = None,
+):
     """Return the indices of the maximum value along the given axes.
 
     Args:
@@ -37,7 +43,7 @@ def argmax_nd(x: np.ndarray, axes: List[int], *, max_rep=np.inf, max_rep_strict=
     if max_rep <= 0:
         raise ValueError("max_rep must be greater than 0.")
     if max_rep_strict is None and not np.isinf(max_rep):
-        raise ValueError("if max_rep_strict is not set if max_rep must be infinite.")
+        raise ValueError("if max_rep_strict is not set, then max_rep must be infinite.")
     # Make it so the axes we want to find the maximum along are the first ones...
     perm = list(range(len(x.shape)))
     for axis in reversed(axes):
@@ -94,14 +100,14 @@ class LayerNMF:
 
     def __init__(
         self,
-        model,
-        layer_name,
-        model_inputs_preprocess,
-        model_inputs_full=None,
+        model: th.nn.Module,
+        layer_name: str,
+        model_inputs_preprocess: th.Tensor,
+        model_inputs_full: Optional[th.Tensor] = None,
         features: Optional[int] = 10,
         *,
         attr_layer_name: Optional[str] = None,
-        attr_opts={"integrate_steps": 10},
+        attr_opts: Dict[str, int] = {"integrate_steps": 10},
         activation_fn: Optional[str] = None,
     ):
         """Use Non-negative matrix factorization dimensionality reduction to then do
@@ -231,9 +237,9 @@ def vis_traditional(
         self,
         feature_list=None,
         *,
-        transforms=[transform.jitter(2)],
-        l2_coeff=0.0,
-        l2_layer_name=None,
+        transforms: List[Callable[[th.Tensor], th.Tensor]] = [transform.jitter(2)],
+        l2_coeff: float = 0.0,
+        l2_layer_name: Optional[str] = None,
     ):
         if feature_list is None:
             # Feature dim is at index 1
@@ -329,7 +335,14 @@ def get_patch(self, obs_index, pos_h, pos_w, *, expand_mult=1):
         slice_w = slice(int(round(left_w)), int(round(right_w)))
         return self.padded_obses[obs_index, :, slice_h, slice_w]
 
-    def vis_dataset(self, feature, *, subdiv_mult=1, expand_mult=1, top_frac=0.1):
+    def vis_dataset(
+        self,
+        feature: Union[int, List[int]],
+        *,
+        subdiv_mult=1,
+        expand_mult=1,
+        top_frac: float = 0.1,
+    ):
         """Visualize a dataset of patches that maximize a given feature.
 
         Args:
@@ -394,7 +407,12 @@ def vis_dataset(self, feature, *, subdiv_mult=1, expand_mult=1, top_frac=0.1):
         )
 
     def vis_dataset_thumbnail(
-        self, feature, *, num_mult=1, expand_mult=1, max_rep=None
+        self,
+        feature: Union[int, List[int]],
+        *,
+        num_mult: int = 1,
+        expand_mult: int = 1,
+        max_rep: Optional[Union[int, float]] = None,
     ):
         """Visualize a dataset of patches that maximize a given feature.