Repository: ad12/meddlr (language: python)

Issue:
v0.0.1a2
## General
- [x] Add MIT license
- [x] Add build, license, pre-commit badges
- [x] Fix setup.py trove classifiers and upload command

## Implementation
### Modeling
- [x] Add block order config parameter for GeneralizedUNet
- [x] Add support for any kind of blocks in unrolled network

### Cluster Management
- [x] Import cluster at the project level to access via `meddlr.Cluster`
- [x] Make cluster instantiation default to current hostname

### config
- [x] Add option to unroll list/tuple/dict arguments when stringifying config params


Discussion / hints:

Fix patch:
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 00000000..6fdc8d5f
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Arjun Desai
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/README.md b/README.md
index 4b447ae3..e98a84ac 100644
--- a/README.md
+++ b/README.md
@@ -1,11 +1,16 @@
 # meddlr
+![GitHub Workflow Status](https://img.shields.io/github/workflow/status/ad12/meddlr/CI)
+![GitHub](https://img.shields.io/github/license/ad12/meddlr)
+[![pre-commit](https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit&logoColor=white)](https://github.com/pre-commit/pre-commit)
+<!-- [![codecov](https://codecov.io/gh/ad12/meddlr/branch/main/graph/badge.svg?token=U6H83UCGFU)](https://codecov.io/gh/ad12/meddlr) -->
+
 [Getting Started](GETTING_STARTED.md)
 
 Meddlr is a config-driven ML framework built to simplify medical image reconstruction and analysis problems.
 
 
 
-#### Installation
+## Installation
 To avoid cuda-related issues, downloading `torch`, `torchvision`, and `cupy` (optional)
 must be done prior to downloading other requirements.
 
@@ -15,25 +20,27 @@ conda create -n meddlr_env python=3.7
 conda activate meddlr_env
 
 # Install cuda-dependant libraries. Change cuda version as needed.
-# Below we show examples for cuda-10.1
-conda install pytorch torchvision cudatoolkit=10.1 -c pytorch
-pip install cupy-cuda101
+# Below we show examples for cuda-10.2
+conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
+pip install cupy-cuda102
+
+# Go to https://github.com/ad12/meddlr and fork the repository.
 
 # Install as package in virtual environment (recommended):
-git clone https://github.com/ad12/meddlr.git
-cd dl-ss-recon && python -m pip install -e '.[dev]'
+git clone https://github.com/<your-github-username>/meddlr.git
+cd meddlr && python -m pip install -e '.[dev]'
 
 # For all contributors, install development packages.
 make dev
 ```
 
-### Contributing
+## Contributing
 See [CONTRIBUTING.md](CONTRIBUTING.md) for more information.
 
-### Acknowledgements
+## Acknowledgements
 Meddlr's design for rapid experimentation and benchmarking is inspired by [detectron2](https://github.com/facebookresearch/detectron2).
 
-### About
+## About
 If you use Meddlr for your work, please consider citing the following work:
 
 ```
diff --git a/codecov.yaml b/codecov.yaml
new file mode 100644
index 00000000..686f86dc
--- /dev/null
+++ b/codecov.yaml
@@ -0,0 +1,5 @@
+coverage:
+  status:
+    # Enable coverage measurement for diff introduced in the pull-request,
+    # but do not mark "X" on commit status for now.
+    patch: off
\ No newline at end of file
diff --git a/meddlr/__init__.py b/meddlr/__init__.py
index b7d20bc4..b3448d0d 100644
--- a/meddlr/__init__.py
+++ b/meddlr/__init__.py
@@ -1,7 +1,6 @@
-from .utils.env import setup_environment
-
-setup_environment()
+# flake8: noqa: F401
 
+from meddlr.utils.cluster import Cluster
 
 # This line will be programatically read/write by setup.py.
 # Leave them at the bottom of this file and don't touch them.
diff --git a/meddlr/config/config.py b/meddlr/config/config.py
index 03cd0def..9c57ac80 100644
--- a/meddlr/config/config.py
+++ b/meddlr/config/config.py
@@ -1,13 +1,10 @@
 """CfgNode implementation.
-
-Adapted from
-https://github.com/facebookresearch/detectron2
 """
 import functools
 import inspect
 import logging
 import re
-from typing import Mapping
+from typing import Any, Mapping
 
 import numpy as np
 from fvcore.common.config import CfgNode as _CfgNode
@@ -28,6 +25,9 @@ class CfgNode(_CfgNode):
 
     # Note that the default value of allow_unsafe is changed to True
     def merge_from_file(self, cfg_filename: str, allow_unsafe: bool = True) -> None:
+        """
+        Adapted from https://github.com/facebookresearch/detectron2
+        """
         loaded_cfg = _CfgNode.load_yaml_with_base(cfg_filename, allow_unsafe=allow_unsafe)
         loaded_cfg = type(self)(loaded_cfg)
 
@@ -69,14 +69,37 @@ def merge_from_file(self, cfg_filename: str, allow_unsafe: bool = True) -> None:
             self.clear()
             self.update(new_config)
 
-    def format_fields(self):
+    def format_fields(self, unroll: bool = False):
         """
         Format string fields in the config by filling them in
         with different parameter values.
+
+        The operation is done in place.
+
+        Args:
+            unroll (bool, optional): If ``True``, sequence types (e.g. tuple, list)
+                stringified as ``seq[0]-seq[1]-...``. Dict types will be
+                stringified as ``k[0]=v[0]-k[1]=v[1]-...``.
+
+        Returns:
+            CfgNode: self
         """
-        return format_config_fields(self, inplace=True)
+        return format_config_fields(self, unroll=unroll, inplace=True)
 
-    def get_recursive(self, key, default=np._NoValue):
+    def get_recursive(self, key, default: Any = np._NoValue):
+        """Get a key recursively from the config.
+
+        Args:
+            key (str): The dot-separated key.
+            default (Any, optional): The value to return if the key is not found.
+                If not provided, a KeyError will be raised.
+
+        Returns:
+            object: The value corresponding to the key.
+
+        Raises:
+            KeyError: If the key is not found and no default was provided.
+        """
         d = self
         try:
             for k in key.split("."):
@@ -87,7 +110,13 @@ def get_recursive(self, key, default=np._NoValue):
             raise e
         return d
 
-    def set_recursive(self, name, value):
+    def set_recursive(self, name: str, value: Any):
+        """Set a key recursively in the config.
+
+        Args:
+            name (str): The dot-separated key.
+            value (object): The value to set.
+        """
         cfg = self
         keys = name.split(".")
         for k in keys[:-1]:
@@ -104,12 +133,20 @@ def dump(self, *args, **kwargs):
         return super().dump(*args, **kwargs)
 
     def freeze(self):
-        """Make this CfgNode and all of its children immutable."""
+        """Make this CfgNode and all of its children immutable.
+
+        Returns:
+            CfgNode: self
+        """
         super().freeze()
         return self
 
     def defrost(self):
-        """Make this CfgNode and all of its children mutable."""
+        """Make this CfgNode and all of its children mutable.
+
+        Returns:
+            CfgNode: self
+        """
         super().defrost()
         return self
 
@@ -170,6 +207,9 @@ def configurable(init_func=None, *, from_config=None):
     Decorate a function or a class's __init__ method so that it can be called
     with a :class:`CfgNode` object using a :func:`from_config` function that translates
     :class:`CfgNode` to arguments.
+
+    Adapted from https://github.com/facebookresearch/detectron2
+
     Examples:
     ::
         # Usage 1: Decorator on __init__:
@@ -308,7 +348,16 @@ def _find_format_str_keys(cfg: Mapping, prefix="", accum=()):
     return accum
 
 
-def format_config_fields(cfg: CfgNode, inplace=False):
+def _unroll_value_to_str(value) -> str:
+    if isinstance(value, (tuple, list)):
+        return "-".join(_unroll_value_to_str(v) for v in value)
+    elif isinstance(value, dict):
+        return "-".join(f"{k}={_unroll_value_to_str(v)}" for k, v in value.items())
+    else:
+        return str(value)
+
+
+def format_config_fields(cfg: CfgNode, unroll=False, inplace=False):
     keys_and_val_str = _find_format_str_keys(cfg)
     values_list = []
     for k, val_str in keys_and_val_str:
@@ -317,6 +366,8 @@ def format_config_fields(cfg: CfgNode, inplace=False):
         assert len(start) == len(end), f"Could not determine formatting string: {val_str}"
         cfg_keys_to_search = [val_str[s + 1 : e] for s, e in zip(start, end)]
         values = [cfg.get_recursive(v) for v in cfg_keys_to_search]
+        if unroll:
+            values = [_unroll_value_to_str(v) for v in values]
 
         fmt_str = ""
         idxs = [0] + [y for x in zip(start, end) for y in x] + [len(val_str)]
diff --git a/meddlr/config/defaults.py b/meddlr/config/defaults.py
index beba9273..c9eaa708 100644
--- a/meddlr/config/defaults.py
+++ b/meddlr/config/defaults.py
@@ -34,6 +34,7 @@
 # Unrolled model - TODO: Deprecate
 # -----------------------------------------------------------------------------
 _C.MODEL.UNROLLED = CN()
+_C.MODEL.UNROLLED.BLOCK_ARCHITECTURE = "ResNet"
 _C.MODEL.UNROLLED.NUM_UNROLLED_STEPS = 5
 _C.MODEL.UNROLLED.NUM_RESBLOCKS = 2
 _C.MODEL.UNROLLED.NUM_FEATURES = 256
@@ -104,7 +105,7 @@
 _C.MODEL.SEG.ACTIVATION = "sigmoid"
 
 # -----------------------------------------------------------------------------
-# UNET model
+# UNET model - shared by UNetModel and GeneralizedUNet modules
 # -----------------------------------------------------------------------------
 _C.MODEL.UNET = CN()
 _C.MODEL.UNET.OUT_CHANNELS = 2
@@ -112,6 +113,9 @@
 _C.MODEL.UNET.CHANNELS = 32
 _C.MODEL.UNET.NUM_POOL_LAYERS = 4
 _C.MODEL.UNET.DROPOUT = 0.0
+_C.MODEL.UNET.NORMALIZE = False
+# BLOCK_ORDER only applies to GeneralizedUNet
+_C.MODEL.UNET.BLOCK_ORDER = ("conv", "relu", "conv", "relu", "batchnorm", "dropout")
 
 # -----------------------------------------------------------------------------
 # Denoising model
diff --git a/meddlr/data/transforms/subsample.py b/meddlr/data/transforms/subsample.py
index f4addc9f..d9852e38 100644
--- a/meddlr/data/transforms/subsample.py
+++ b/meddlr/data/transforms/subsample.py
@@ -2,6 +2,7 @@
 import os
 from typing import Sequence
 
+import numba as nb
 import numpy as np
 import sigpy.mri
 import torch
@@ -120,15 +121,19 @@ def __init__(
         center_fractions=None,
         max_attempts=30,
         crop_corner=True,
+        module="internal",
     ):
         if center_fractions:
             raise ValueError(f"center_fractions not yet supported for class {type(self)}.")
+        if module not in ("internal", "sigpy"):
+            raise ValueError("`module` must be one of ('internal', 'sigpy')")
         super().__init__(accelerations)
         if isinstance(calib_size, int):
             calib_size = (calib_size, calib_size)
         self.calib_size = calib_size
         self.max_attempts = max_attempts
         self.crop_corner = crop_corner
+        self.module = module
 
     def __call__(self, out_shape, seed=None, acceleration=None):
         # Design parameters for mask
@@ -146,15 +151,30 @@ def __call__(self, out_shape, seed=None, acceleration=None):
             shape = (nky, nkz)
             transpose = False
 
-        mask = sigpy.mri.poisson(
-            shape,
-            acceleration,
-            calib=self.calib_size,
-            dtype=np.float32,
-            seed=seed,
-            max_attempts=self.max_attempts,
-            crop_corner=self.crop_corner,
-        )
+        # Issue #2: Due to some optimization reasons, using the internal
+        # poisson disc module has been faster than using sigpy's
+        # default one. In many cases, the call to sigpy hangs.
+        module = self.module
+        if module == "internal":
+            mask = poisson(
+                shape,
+                acceleration,
+                calib=self.calib_size,
+                dtype=np.float32,
+                seed=seed,
+                K=self.max_attempts,
+                crop_corner=self.crop_corner,
+            )
+        elif module == "sigpy":
+            mask = sigpy.mri.poisson(
+                shape,
+                acceleration,
+                calib=self.calib_size,
+                dtype=np.float32,
+                seed=seed,
+                max_attempts=self.max_attempts,
+                crop_corner=self.crop_corner,
+            )
         if transpose:
             mask = mask.transpose()
 
@@ -306,3 +326,151 @@ def __call__(self, out_shape, seed=None, acceleration=None):
 
         mask = mask.reshape(out_shape)
         return torch.from_numpy(mask)
+
+
+# ================================================================ #
+# Adapted from sigpy.
+# Duplicated because of https://github.com/mikgroup/sigpy/issues/54
+# TODO: Remove once https://github.com/mikgroup/sigpy/issues/54 is
+# solved and added to release.
+# ================================================================ #
+def poisson(
+    img_shape,
+    accel,
+    K=30,
+    calib=(0, 0),
+    dtype=np.complex,
+    crop_corner=True,
+    return_density=False,
+    seed=0,
+):
+    """Generate Poisson-disc sampling pattern
+
+    Args:
+        img_shape (tuple of ints): length-2 image shape.
+        accel (float): Target acceleration factor. Greater than 1.
+        K (float): maximum number of samples to reject.
+        calib (tuple of ints): length-2 calibration shape.
+        dtype (Dtype): data type.
+        crop_corner (bool): Toggle whether to crop sampling corners.
+        return_density (bool): Toggle whether to return sampling density.
+        seed (int): Random seed.
+
+    Returns:
+        array: Poisson-disc sampling mask.
+
+    References:
+        Bridson, Robert. "Fast Poisson disk sampling in arbitrary dimensions."
+        SIGGRAPH sketches. 2007.
+
+    """
+    y, x = np.mgrid[: img_shape[-2], : img_shape[-1]]
+    x = np.maximum(abs(x - img_shape[-1] / 2) - calib[-1] / 2, 0)
+    x /= x.max()
+    y = np.maximum(abs(y - img_shape[-2] / 2) - calib[-2] / 2, 0)
+    y /= y.max()
+    r = np.sqrt(x ** 2 + y ** 2)
+
+    slope_max = 40
+    slope_min = 0
+    if seed is not None:
+        rand_state = np.random.get_state()
+    else:
+        seed = -1  # numba does not play nicely with None types
+    while slope_min < slope_max:
+        slope = (slope_max + slope_min) / 2.0
+        R = 1.0 + r * slope
+        mask = _poisson(img_shape[-1], img_shape[-2], K, R, calib, seed)
+        if crop_corner:
+            mask *= r < 1
+
+        est_accel = img_shape[-1] * img_shape[-2] / np.sum(mask[:])
+
+        if abs(est_accel - accel) < 0.1:
+            break
+        if est_accel < accel:
+            slope_min = slope
+        else:
+            slope_max = slope
+
+    if seed is not None and seed > 0:
+        np.random.set_state(rand_state)
+    mask = mask.reshape(img_shape).astype(dtype)
+    if return_density:
+        return mask, r
+    else:
+        return mask
+
+
+@nb.jit(nopython=True, cache=True)  # pragma: no cover
+def _poisson(nx, ny, K, R, calib, seed=None):
+
+    mask = np.zeros((ny, nx))
+    f = ny / nx
+
+    if seed is not None and seed > 0:
+        np.random.seed(int(seed))
+
+    pxs = np.empty(nx * ny, np.int32)
+    pys = np.empty(nx * ny, np.int32)
+    pxs[0] = np.random.randint(0, nx)
+    pys[0] = np.random.randint(0, ny)
+    m = 1
+    while m > 0:
+
+        i = np.random.randint(0, m)
+        px = pxs[i]
+        py = pys[i]
+        rad = R[py, px]
+
+        # Attempt to generate point
+        done = False
+        k = 0
+        while not done and k < K:
+
+            # Generate point randomly from R and 2R
+            rd = rad * (np.random.random() * 3 + 1) ** 0.5
+            t = 2 * np.pi * np.random.random()
+            qx = px + rd * np.cos(t)
+            qy = py + rd * f * np.sin(t)
+
+            # Reject if outside grid or close to other points
+            if qx >= 0 and qx < nx and qy >= 0 and qy < ny:
+
+                startx = max(int(qx - rad), 0)
+                endx = min(int(qx + rad + 1), nx)
+                starty = max(int(qy - rad * f), 0)
+                endy = min(int(qy + rad * f + 1), ny)
+
+                done = True
+                for x in range(startx, endx):
+                    for y in range(starty, endy):
+                        if mask[y, x] == 1 and (
+                            ((qx - x) / R[y, x]) ** 2 + ((qy - y) / (R[y, x] * f)) ** 2 < 1
+                        ):
+                            done = False
+                            break
+
+                    if not done:
+                        break
+
+            k += 1
+
+        # Add point if done else remove active
+        if done:
+            pxs[m] = qx
+            pys[m] = qy
+            mask[int(qy), int(qx)] = 1
+            m += 1
+        else:
+            pxs[i] = pxs[m - 1]
+            pys[i] = pys[m - 1]
+            m -= 1
+
+    # Add calibration region
+    mask[
+        int(ny / 2 - calib[-2] / 2) : int(ny / 2 + calib[-2] / 2),
+        int(nx / 2 - calib[-1] / 2) : int(nx / 2 + calib[-1] / 2),
+    ] = 1
+
+    return mask
diff --git a/meddlr/modeling/meta_arch/generalized_unet.py b/meddlr/modeling/meta_arch/generalized_unet.py
index 39004f63..15956a75 100644
--- a/meddlr/modeling/meta_arch/generalized_unet.py
+++ b/meddlr/modeling/meta_arch/generalized_unet.py
@@ -163,5 +163,8 @@ def from_config(cls, cfg, **kwargs):
             "channels": num_channels,
             "dropout": cfg.DROPOUT,
         }
+        block_order = cfg.get("BLOCK_ORDER", None)
+        if block_order is not None:
+            init_args["block_order"] = block_order
         init_args.update(**kwargs)
         return init_args
diff --git a/meddlr/modeling/meta_arch/unet.py b/meddlr/modeling/meta_arch/unet.py
index 4ecfe585..68acae05 100644
--- a/meddlr/modeling/meta_arch/unet.py
+++ b/meddlr/modeling/meta_arch/unet.py
@@ -122,6 +122,7 @@ def __init__(
         dropout: float = 0.0,
         use_latent: bool = False,
         num_latent_layers: int = 1,
+        normalize: bool = False,
         vis_period: int = -1,
     ):
         """
@@ -131,6 +132,7 @@ def __init__(
             chans (int): Number of output channels of the first convolution layer.
             num_pool_layers (int): Number of down-sampling and up-sampling layers.
             drop_prob (float): Dropout probability.
+            normalize (bool, optional): If ``True``, normalize the input for stability.
         """
         super().__init__()
 
@@ -141,6 +143,7 @@ def __init__(
         self.drop_prob = dropout
         self.use_latent = use_latent
         self.num_latent_layers = num_latent_layers
+        self.normalize = normalize
 
         self.down_sample_layers = nn.ModuleList([ConvBlock(in_channels, channels, dropout)])
         ch = channels
@@ -226,6 +229,70 @@ def visualize_training(self, kspace, zfs, targets, preds):
                 data = tv_utils.make_grid(data, nrow=1, padding=1, normalize=True, scale_each=True)
                 storage.put_image("train/{}".format(name), data.numpy(), data_format="CHW")
 
+    def norm(self, x):
+        """
+        Args:
+            x (torch.Tensor): Input tensor. Shape [B, 2*C, H, W, ...]
+        """
+        ndim = x.ndim
+        x = x.view(x.shape[:2] + (-1,))
+
+        mean = x.mean(dim=-1).reshape(x.shape[:2] + (1,) * (ndim - 2))
+        std = x.std(dim=-1).reshape(x.shape[:2] + (1,) * (ndim - 2))
+
+        return (x - mean) / std, mean, std
+
+    def unnorm(self, x, mean, std):
+        return x * std + mean
+
+    def base_forward(self, zf_image):
+        """The base forward function.
+
+        Args:
+            zf_image (torch.Tensor): Zero-filled image. Shape [B, C, ..., H, W].
+                The real and imaginary channels should be part of the channels
+                dimension if the input is complex and the model does not have
+                complex weights.
+
+        Returns:
+            torch.Tensor: Reconstructed image. Shape [B, C, ..., H, W].
+        """
+        x = zf_image
+
+        stack = []
+        if self.normalize:
+            x, mean, std = self.norm(x)
+
+        # Apply down-sampling layers
+        for layer in self.down_sample_layers:
+            x = layer(x)
+            stack.append(x)
+            x = F.avg_pool2d(x, kernel_size=2, stride=2, padding=0)
+
+        x = self.conv(x)
+
+        # Apply up-sampling layers
+        for transpose_conv, conv in zip(self.up_transpose_conv, self.up_conv):
+            downsample_layer = stack.pop()
+            x = transpose_conv(x)
+
+            # Reflect pad on the right/botton if needed to handle odd input dimensions.
+            padding = [0, 0, 0, 0]
+            if x.shape[-1] != downsample_layer.shape[-1]:
+                padding[1] = 1  # Padding right
+            if x.shape[-2] != downsample_layer.shape[-2]:
+                padding[3] = 1  # Padding bottom
+            if sum(padding) != 0:
+                x = F.pad(x, padding, "reflect")
+
+            x = torch.cat([x, downsample_layer], dim=1)
+            x = conv(x)
+
+        if self.normalize:
+            x = self.unnorm(x, mean, std)
+
+        return x
+
     def forward(self, input, return_pp=False, vis_training=False):
         """
         Args:
@@ -234,7 +301,6 @@ def forward(self, input, return_pp=False, vis_training=False):
             (torch.Tensor): Output tensor of shape [batch_size, self.out_chans, height, width]
         """
         self.register_hooks()
-        stack = []
 
         # Need to fetch device at runtime for proper data transfer.
         inputs = input
@@ -244,11 +310,6 @@ def forward(self, input, return_pp=False, vis_training=False):
         mask = inputs["mask"].to(device) if "mask" in inputs else None
         A = inputs["signal_model"].to(device) if "signal_model" in inputs else None
         maps = inputs["maps"].to(device)
-        # num_maps_dim = -2 if cplx.is_complex_as_real(maps) else -1
-        # if self.num_emaps != maps.size()[num_maps_dim]:
-        #     raise ValueError(
-        #         "Incorrect number of ESPIRiT maps! Re-prep data..."
-        #     )
 
         if mask is None:
             mask = cplx.get_mask(kspace)
@@ -278,32 +339,9 @@ def forward(self, input, return_pp=False, vis_training=False):
             output = zf_image
         output = zf_image.permute(0, 4, 1, 2, 3).squeeze(-1)
 
-        # Apply down-sampling layers
-        for layer in self.down_sample_layers:
-            output = layer(output)
-            stack.append(output)
-            output = F.avg_pool2d(output, kernel_size=2, stride=2, padding=0)
-
-        output = self.conv(output)
-
-        # Apply up-sampling layers
-        for transpose_conv, conv in zip(self.up_transpose_conv, self.up_conv):
-            downsample_layer = stack.pop()
-            output = transpose_conv(output)
-
-            # Reflect pad on the right/botton if needed to handle odd input dimensions.
-            padding = [0, 0, 0, 0]
-            if output.shape[-1] != downsample_layer.shape[-1]:
-                padding[1] = 1  # Padding right
-            if output.shape[-2] != downsample_layer.shape[-2]:
-                padding[3] = 1  # Padding bottom
-            if sum(padding) != 0:
-                output = F.pad(output, padding, "reflect")
-
-            output = torch.cat([output, downsample_layer], dim=1)
-            output = conv(output)
+        # Run U-Net.
+        output = self.base_forward(output)
 
-        # pred = output.view(zf_dims)
         if num_maps > 1:
             pred = output.permute(0, 2, 3, 1)
             pred = pred.reshape(pred.shape[:-1] + (num_maps, 2))
@@ -344,5 +382,6 @@ def from_config(cls, cfg):
             "dropout": cfg.MODEL.UNET.DROPOUT,
             "use_latent": cfg.get_recursive("MODEL.CONSISTENCY.USE_LATENT", False),
             "num_latent_layers": cfg.get_recursive("MODEL.CONSISTENCY.NUM_LATENT_LAYERS", 1),
+            "normalize": cfg.MODEL.UNET.NORMALIZE,
             "vis_period": cfg.VIS_PERIOD,
         }
diff --git a/meddlr/modeling/meta_arch/unrolled.py b/meddlr/modeling/meta_arch/unrolled.py
index 1cd4bd02..fdda4760 100644
--- a/meddlr/modeling/meta_arch/unrolled.py
+++ b/meddlr/modeling/meta_arch/unrolled.py
@@ -12,7 +12,7 @@
 from meddlr.utils.transforms import SenseModel
 
 from ..layers.layers2D import ResNet
-from .build import META_ARCH_REGISTRY
+from .build import META_ARCH_REGISTRY, build_model
 
 __all__ = ["GeneralizedUnrolledCNN"]
 
@@ -126,7 +126,7 @@ def forward(self, inputs, return_pp=False, vis_training=False):
         if vis_training and not self.training:
             raise ValueError("vis_training is only applicable in training mode.")
         # Need to fetch device at runtime for proper data transfer.
-        device = self.resnets[0].final_layer.weight.device
+        device = next(self.resnets[0].parameters()).device
         inputs = move_to_device(inputs, device)
         kspace = inputs["kspace"]
         target = inputs.get("target", None)
@@ -173,9 +173,14 @@ def forward(self, inputs, return_pp=False, vis_training=False):
 
             # prox update
             image = image.reshape(dims[0:3] + (self.num_emaps * 2,)).permute(0, 3, 1, 2)
-            image = resnet(image)
+            if hasattr(resnet, "base_forward") and callable(resnet.base_forward):
+                image = resnet.base_forward(image)
+            else:
+                image = resnet(image)
 
             image = image.permute(0, 2, 3, 1).reshape(dims[0:3] + (self.num_emaps, 2))
+            if not image.is_contiguous():
+                image = image.contiguous()
             if use_cplx:
                 image = torch.view_as_complex(image)
 
@@ -202,37 +207,27 @@ def from_config(cls, cfg):
         """
         # Extract network parameters
         num_grad_steps = cfg.MODEL.UNROLLED.NUM_UNROLLED_STEPS
-        num_resblocks = cfg.MODEL.UNROLLED.NUM_RESBLOCKS
-        num_features = cfg.MODEL.UNROLLED.NUM_FEATURES
-        kernel_size = cfg.MODEL.UNROLLED.KERNEL_SIZE
-        if len(kernel_size) == 1:
-            kernel_size = kernel_size[0]
-        drop_prob = cfg.MODEL.UNROLLED.DROPOUT
-        circular_pad = cfg.MODEL.UNROLLED.PADDING == "circular"
         share_weights = cfg.MODEL.UNROLLED.SHARE_WEIGHTS
 
         # Data dimensions
         num_emaps = cfg.MODEL.UNROLLED.NUM_EMAPS
 
-        # ResNet parameters
-        resnet_params = dict(
-            num_resblocks=num_resblocks,
-            in_chans=2 * num_emaps,
-            chans=num_features,
-            kernel_size=kernel_size,
-            drop_prob=drop_prob,
-            circular_pad=circular_pad,
-            act_type=cfg.MODEL.UNROLLED.CONV_BLOCK.ACTIVATION,
-            norm_type=cfg.MODEL.UNROLLED.CONV_BLOCK.NORM,
-            norm_affine=cfg.MODEL.UNROLLED.CONV_BLOCK.NORM_AFFINE,
-            order=cfg.MODEL.UNROLLED.CONV_BLOCK.ORDER,
-        )
+        # Determine block to use for each unrolled step.
+        if cfg.MODEL.UNROLLED.BLOCK_ARCHITECTURE == "ResNet":
+            builder = lambda: _build_resblock(cfg)  # noqa: E731
+        else:
+            # TODO: Fix any inconsistencies between config's IN_CHANNELS
+            # and the number of channels that the unrolled net expects.
+            mcfg = cfg.clone().defrost()
+            mcfg.MODEL.META_ARCHITECTURE = cfg.MODEL.UNROLLED.BLOCK_ARCHITECTURE
+            mcfg = mcfg.freeze()
+            builder = lambda: build_model(mcfg)  # noqa: E731
 
         # Declare ResNets and RNNs for each unrolled iteration
         if share_weights:
-            blocks = nn.ModuleList([ResNet(**resnet_params)] * num_grad_steps)
+            blocks = nn.ModuleList([builder()] * num_grad_steps)
         else:
-            blocks = nn.ModuleList([ResNet(**resnet_params) for _ in range(num_grad_steps)])
+            blocks = nn.ModuleList([builder() for _ in range(num_grad_steps)])
 
         return {
             "blocks": blocks,
@@ -240,3 +235,37 @@ def from_config(cls, cfg):
             "num_emaps": num_emaps,
             "vis_period": cfg.VIS_PERIOD,
         }
+
+
+def _build_resblock(cfg):
+    """Build the resblock for unrolled network.
+
+    Args:
+        cfg (CfgNode): The network configuration.
+
+    Note:
+        This is a temporary method used as a base case for building
+        unrolled networks with the default resblocks. In the future,
+        this will be handled by :func:`meddlr.modeling.meta_arch.build_model`.
+    """
+    # Data dimensions
+    num_emaps = cfg.MODEL.UNROLLED.NUM_EMAPS
+
+    # ResNet parameters
+    kernel_size = cfg.MODEL.UNROLLED.KERNEL_SIZE
+    if len(kernel_size) == 1:
+        kernel_size = kernel_size[0]
+    resnet_params = dict(
+        num_resblocks=cfg.MODEL.UNROLLED.NUM_RESBLOCKS,
+        in_chans=2 * num_emaps,  # complex -> real/imag
+        chans=cfg.MODEL.UNROLLED.NUM_FEATURES,
+        kernel_size=kernel_size,
+        drop_prob=cfg.MODEL.UNROLLED.DROPOUT,
+        circular_pad=cfg.MODEL.UNROLLED.PADDING == "circular",
+        act_type=cfg.MODEL.UNROLLED.CONV_BLOCK.ACTIVATION,
+        norm_type=cfg.MODEL.UNROLLED.CONV_BLOCK.NORM,
+        norm_affine=cfg.MODEL.UNROLLED.CONV_BLOCK.NORM_AFFINE,
+        order=cfg.MODEL.UNROLLED.CONV_BLOCK.ORDER,
+    )
+
+    return ResNet(**resnet_params)
diff --git a/meddlr/utils/cluster.py b/meddlr/utils/cluster.py
index b4dd61ec..9e507391 100644
--- a/meddlr/utils/cluster.py
+++ b/meddlr/utils/cluster.py
@@ -17,6 +17,8 @@
 # TODO: make this cleaner
 _REPO_DIR = os.path.join(os.path.dirname(__file__), "../..")
 
+__all__ = ["Cluster", "set_cluster"]
+
 
 class Cluster:
     """Manages configurations for different nodes/clusters.
@@ -38,7 +40,7 @@ class Cluster:
 
     To use the configurations of a particular cluster, set the working cluster:
 
-    >>> Cluster.set_working_cluster(cluster)
+    >>> cluster.use()
 
     Configs can be persisted by saving to a file. If the config has been saved,
     future sessions will attempt to auto-detect the set of saved configs:
@@ -65,8 +67,8 @@ class Cluster:
 
     def __init__(
         self,
-        name: str,
-        patterns: Union[str, Sequence[str]],
+        name: str = None,
+        patterns: Union[str, Sequence[str]] = None,
         data_dir: str = None,
         results_dir: str = None,
         cache_dir: str = None,
@@ -78,6 +80,7 @@ def __init__(
             patterns (Sequence[str]): Regex pattern(s) for identifying nodes
                 in the cluster. Cluster will be identified by
                 ``any(re.match(p, socket.gethostname()) for p in patterns)``.
+                If ``None``, defaults to current hostname.
             data_dir (str, optional): The data directory. Defaults to
                 ``os.environ['MEDDLR_DATASETS_DIR']`` or ``"./datasets"``.
             results_dir (str, optional): The results directory. Defaults to
@@ -85,10 +88,14 @@ def __init__(
             cfg_kwargs (optional): Any other configurations you would like to
                 store for the cluster.
         """
+        if name is None:
+            name = socket.gethostname()
         self.name = name
 
+        if patterns is None:
+            patterns = socket.gethostname()
         if isinstance(patterns, str):
-            patterns = patterns
+            patterns = [patterns]
         self.patterns = patterns
 
         self._data_dir = data_dir
@@ -102,6 +109,11 @@ def data_dir(self):
         path = os.environ.get("MEDDLR_DATASETS_DIR", path if path else "./datasets")
         return PathManager.get_local_path(path)
 
+    @property
+    def datasets_dir(self):
+        """Alias for ``self.data_dir``."""
+        return self.data_dir
+
     @property
     def results_dir(self):
         path = self._results_dir
@@ -111,13 +123,34 @@ def results_dir(self):
     @property
     def cache_dir(self):
         path = self._cache_dir
-        path = os.environ.get("MEDDLR_CACHE_DIR", path if path else "~/cache/meddl")
+        path = os.environ.get("MEDDLR_CACHE_DIR", path if path else "~/cache/meddlr")
         return PathManager.get_local_path(path)
 
+    def set(self, **kwargs):
+        """Set cluster configuration properties.
+
+        Args:
+            kwargs: Keyword arguments to set.
+
+        Examples:
+            >>> cluster.set(data_dir="/path/to/datasets", results_dir="/path/to/results")
+
+        Note:
+            Setting attributes for the cluster will not override environment variables.
+            For example, if ``os.environ['MEDDLR_DATASETS_DIR']`` is set, setting ``data_dir``
+            will have no effect.
+        """
+        for k, v in kwargs.items():
+            private_key = f"_{k}"
+            if hasattr(self, private_key):
+                setattr(self, private_key, v)
+            else:
+                self._cfg_kwargs[k] = v
+
     def __getattr__(self, attr: str):
         attr_env = f"MEDDLR_{attr.upper()}"
         try:
-            value = self._cfg_kwargs.get(attr, os.environ[attr_env])
+            value = os.environ.get(attr_env, self._cfg_kwargs[attr])
         except KeyError:
             raise AttributeError(f"Attribute {attr} not specified for cluster {self.name}.")
         return value
@@ -217,15 +250,20 @@ def config_file():
     def working_cluster() -> "Cluster":
         return _CLUSTER
 
-    @staticmethod
-    def set_working_cluster(cluster=None):
-        """Sets the working cluster.
-        Args:
-            cluster (`str` or `Cluster`): The cluster name or cluster.
-                If ``None``, will reset cluster to _UNKNOWN, meaning default
-                data and results dirs will be used.
+    def use(self):
+        """Sets ``self`` to be the working cluster of the project.
+
+        The working cluster is the default cluster that is used to manage
+        paths and other configuration variables.
+
+        Examples:
+            >>> cluster.use()
+
+        Note:
+            This function does not override environment variables.
+            All environment variables will take priority over this clusters
         """
-        set_cluster(cluster)
+        set_cluster(self)
 
     def __repr__(self):
         return "Cluster({})".format(


Test patch:
diff --git a/configs/tests/unrolled-unet.yaml b/configs/tests/unrolled-unet.yaml
new file mode 100644
index 00000000..4eb55c3e
--- /dev/null
+++ b/configs/tests/unrolled-unet.yaml
@@ -0,0 +1,17 @@
+# Using the U-Net as blocks in the unrolled network.
+_BASE_: "basic.yaml"
+MODEL:
+  META_ARCHITECTURE: "GeneralizedUnrolledCNN"
+  UNROLLED:
+    BLOCK_ARCHITECTURE: "UnetModel"
+  UNET:
+    IN_CHANNELS: 2
+    OUT_CHANNELS: 2
+    CHANNELS: 32
+    NUM_POOL_LAYERS: 4
+    DROPOUT: 0.
+  RECON_LOSS:
+    NAME: "l1"
+    RENORMALIZE_DATA: False
+OUTPUT_DIR: "results://tests/unrolled-unet"
+VERSION: 1
\ No newline at end of file
diff --git a/tests/config/test_config.py b/tests/config/test_config.py
index df55c014..59e91d4e 100644
--- a/tests/config/test_config.py
+++ b/tests/config/test_config.py
@@ -19,3 +19,15 @@ def test_format_fields(self):
         cfg.DESCRIPTION.BRIEF = 'f"{SEED},project={DESCRIPTION.PROJECT_NAME}-today"'
         cfg.format_fields()
         assert cfg.DESCRIPTION.BRIEF == f"{cfg.SEED},project={cfg.DESCRIPTION.PROJECT_NAME}-today"
+
+        cfg = get_cfg()
+        cfg.AUG_TRAIN.UNDERSAMPLE.ACCELERATIONS = (6, 8)
+        cfg.DESCRIPTION.BRIEF = (
+            'f"{AUG_TRAIN.UNDERSAMPLE.ACCELERATIONS},project={DESCRIPTION.PROJECT_NAME}-today"'
+        )
+        cfg.format_fields(unroll=True)
+        expected = "{},project={}-today".format(
+            "-".join(str(x) for x in cfg.AUG_TRAIN.UNDERSAMPLE.ACCELERATIONS),
+            cfg.DESCRIPTION.PROJECT_NAME,
+        )
+        assert cfg.DESCRIPTION.BRIEF == expected
diff --git a/tests/data/transforms/__init__.py b/tests/data/transforms/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/data/transforms/test_subsample.py b/tests/data/transforms/test_subsample.py
new file mode 100644
index 00000000..7426a7d1
--- /dev/null
+++ b/tests/data/transforms/test_subsample.py
@@ -0,0 +1,20 @@
+import numpy as np
+import torch
+
+from meddlr.data.transforms.subsample import PoissonDiskMaskFunc
+
+
+def test_poisson_disc():
+    """
+    Verify that built-in Poisson Disc generator returns same value as SigPy.
+    """
+    acc = 6
+    calib_size = 20
+    shape = (1, 320, 256)  # (batch, height, width)
+    seed = 49
+
+    pd_builtin = PoissonDiskMaskFunc(acc, calib_size)
+    builtin_mask = pd_builtin(shape, seed=seed, acceleration=acc)
+
+    # assert all close with 5% tolerance.
+    assert np.allclose(torch.sum(builtin_mask).item() / np.prod(shape), 1 / acc, atol=5e-2)
diff --git a/tests/util.py b/tests/util.py
new file mode 100644
index 00000000..b922852d
--- /dev/null
+++ b/tests/util.py
@@ -0,0 +1,19 @@
+"""Utilities for testing."""
+import os
+from functools import wraps
+
+
+def temp_env(func):
+    """Allows func to temporarily set environment variables."""
+
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        old_env = dict(os.environ)
+
+        out = func(self, *args, **kwargs)
+
+        os.environ.clear()
+        os.environ.update(old_env)
+        return out
+
+    return wrapper
diff --git a/tests/utils/test_cluster.py b/tests/utils/test_cluster.py
new file mode 100644
index 00000000..ddda8b5d
--- /dev/null
+++ b/tests/utils/test_cluster.py
@@ -0,0 +1,75 @@
+import os
+import socket
+import unittest
+
+from meddlr.utils.cluster import Cluster
+
+from ..util import temp_env
+
+
+class TestCluster(unittest.TestCase):
+    def test_basic(self):
+        """Test basic cluster creation."""
+        cluster = Cluster()
+
+        assert cluster.name == socket.gethostname()
+        assert cluster.patterns == [socket.gethostname()]
+
+        assert cluster._data_dir is None
+        assert cluster._results_dir is None
+        assert cluster._cache_dir is None
+        assert cluster._cfg_kwargs == {}
+
+        assert cluster.data_dir == "./datasets"
+        assert cluster.datasets_dir == cluster.data_dir
+        assert cluster.results_dir == "./results"
+        assert cluster.cache_dir == "~/cache/meddlr"
+
+    def test_set(self):
+        """Test setting configuration properties."""
+        cluster = Cluster()
+
+        cluster.set(data_dir="my-data", results_dir="my-results", cache_dir="my-cache")
+        assert cluster.data_dir == "my-data"
+        assert cluster.results_dir == "my-results"
+        assert cluster.cache_dir == "my-cache"
+
+        cluster.set(foo="bar")
+        assert cluster.foo == "bar"
+
+    @temp_env
+    def test_environment_variables(self):
+        """Test that environment variables are overrides for cluster configurations."""
+        cluster = Cluster()
+
+        cluster.set(data_dir="my-data", results_dir="my-results", cache_dir="my-cache")
+        assert cluster.data_dir == "my-data"
+        assert cluster.results_dir == "my-results"
+        assert cluster.cache_dir == "my-cache"
+
+        os.environ["MEDDLR_DATASETS_DIR"] = "foo-bar"
+        assert cluster.data_dir == "foo-bar"
+
+    def test_peristence(self):
+        """Test that cluster configuration can be persisted (saved, loaded, deleted).
+
+        Note:
+            This test is not thread-safe. Please do not run multiple instances of
+            this test simultaneously with multiprocessing.
+        """
+        name = "test-foobar"
+        cluster = Cluster(name=name)
+
+        cluster.set(data_dir="my-data", results_dir="my-results", cache_dir="my-cache")
+
+        cluster.save()
+        cluster2 = cluster.from_config(name)
+        assert cluster2 is not cluster
+        assert cluster2.data_dir == "my-data"
+        assert cluster2.results_dir == "my-results"
+        assert cluster2.cache_dir == "my-cache"
+
+        cluster.delete()
+        with self.assertRaises(KeyError):
+            # Once deleted the cluster should not be found.
+            cluster2 = cluster.from_config(name)