Fix CNN training on singleton vector batches

docs/api/models/pyhealth.models.CNN.rst

@@ -3,6 +3,9 @@
 
 The separate callable CNNLayer and the complete CNN model.
 
+For 1D inputs with one sample and one sequence position, batch normalization
+uses running statistics while preserving gradients and the layers' training state.
+
 .. autoclass:: pyhealth.models.CNNLayer
     :members:
     :undoc-members:

examples/cnn_mimic4.ipynb

@@ -251,7 +251,9 @@
    "metadata": {},
    "source": [
     "# 5. Instantiate CNN Model\n",
-    "Create the PyHealth CNN with custom hyperparameters and inspect the parameter footprint prior to optimisation."
+    "Create the PyHealth CNN with custom hyperparameters and inspect the parameter footprint prior to optimisation.\n",
+    "\n",
+    "If a final batch contains one patient with a length-1 feature, CNN uses BatchNorm's running statistics for that feature while keeping gradients enabled."
    ]
   },
   {

pyhealth/models/cnn.py

@@ -31,6 +31,13 @@ class CNNBlock(nn.Module):
     Args:
         in_channels: number of input channels.
         out_channels: number of output channels.
+
+    Example:
+        >>> import torch
+        >>> from pyhealth.models.cnn import CNNBlock
+        >>> block = CNNBlock(4, 8, spatial_dim=1)
+        >>> block(torch.randn(1, 4, 1)).shape
+        torch.Size([1, 8, 1])
     """
 
     def __init__(self, in_channels: int, out_channels: int, spatial_dim: int):
@@ -61,6 +68,23 @@ def __init__(self, in_channels: int, out_channels: int, spatial_dim: int):
             )
         self.relu = nn.ReLU()
 
+    @staticmethod
+    def _apply_layers(layers: nn.Sequential, x: torch.Tensor) -> torch.Tensor:
+        for layer in layers:
+            if (
+                isinstance(layer, nn.BatchNorm1d)
+                and layer.training
+                and x.size(0) * x.size(2) == 1
+            ):
+                # Use running statistics without changing the layer's mode.
+                x = nn.functional.batch_norm(
+                    x, layer.running_mean, layer.running_var,
+                    layer.weight, layer.bias, training=False, eps=layer.eps,
+                )
+            else:
+                x = layer(x)
+        return x
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward propagation.
 
@@ -71,10 +95,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             output tensor of shape [batch size, out_channels, *].
         """
         residual = x
-        out = self.conv1(x)
-        out = self.conv2(out)
+        out = self._apply_layers(self.conv1, x)
+        out = self._apply_layers(self.conv2, out)
         if self.downsample is not None:
-            residual = self.downsample(x)
+            residual = self._apply_layers(self.downsample, x)
         out += residual
         out = self.relu(out)
         return out

tests/core/test_cnn.py

@@ -277,15 +277,22 @@ def test_model_with_multihot_and_1d_tensor_inputs(self):
         self.assertEqual(model.feature_conv_dims["demographics"], 1)
         self.assertEqual(model.feature_conv_dims["vitals"], 1)
 
-        train_loader = get_dataloader(dataset, batch_size=2, shuffle=False)
+        train_loader = get_dataloader(dataset, batch_size=1, shuffle=False)
         data_batch = next(iter(train_loader))
 
         ret = model(**data_batch)
         ret["loss"].backward()
 
-        self.assertEqual(ret["y_prob"].shape[0], 2)
-        self.assertEqual(ret["logit"].shape[0], 2)
+        self.assertEqual(ret["y_prob"].shape[0], 1)
+        self.assertEqual(ret["logit"].shape[0], 1)
         self.assertEqual(ret["loss"].dim(), 0)
+        self.assertTrue(torch.isfinite(ret["loss"]))
+        for module in model.modules():
+            if isinstance(module, torch.nn.BatchNorm1d):
+                self.assertTrue(module.training)
+                self.assertEqual(module.num_batches_tracked.item(), 0)
+                self.assertIsNotNone(module.weight.grad)
+                self.assertTrue(torch.isfinite(module.weight.grad).all())
 
 
 if __name__ == "__main__":