add stride to patchify and add stride arg to field dataset

rctn · Dec 19, 2023 · 63502c8 · 63502c8
1 parent 64fe099
commit 63502c8
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/sparsecoding/data/datasets/field.py b/sparsecoding/data/datasets/field.py
@@ -21,6 +21,9 @@ class FieldDataset(Dataset):
         Location to download the dataset to.
     patch_size : int
         Side length of patches for sparse dictionary learning.
+    stride : int, optional
+        Stride for sampling patches. If not specified, set to `patch_size` 
+        (non-overlapping patches).
     """
 
     B = 10
@@ -32,8 +35,11 @@ def __init__(
         self,
         root: str,
         patch_size: int = 8,
+        stride: int = None,
     ):
         self.P = patch_size
+        if stride is None:
+            stride = patch_size
 
         root = os.path.expanduser(root)
         os.system(f"mkdir -p {root}")
@@ -47,7 +53,7 @@ def __init__(
         self.images = torch.permute(self.images, (2, 0, 1))  # [B, H, W]
         self.images = torch.reshape(self.images, (self.B, self.C, self.H, self.W))  # [B, C, H, W]
 
-        self.patches = patchify(patch_size, self.images)  # [B, N, C, P, P]
+        self.patches = patchify(patch_size, self.images, stride)  # [B, N, C, P, P]
         self.patches = torch.reshape(self.patches, (-1, self.C, self.P, self.P))  # [B*N, C, P, P]
 
     def __len__(self):

diff --git a/sparsecoding/data/transforms/patch.py b/sparsecoding/data/transforms/patch.py
@@ -65,6 +65,7 @@ def sample_random_patches(
 def patchify(
     patch_size: int,
     image: torch.Tensor,
+    stride: int = None,
 ):
     """Break an image into square patches.
 
@@ -79,6 +80,9 @@ def patchify(
             C is the number of channels,
             H is the image height,
             W is the image width.
+    stride : int, optional 
+        Stride between patches in pixel space. If not specified, set to 
+        `patch_size` (non-overlapping patches).
 
     Returns
     -------
@@ -92,6 +96,8 @@ def patchify(
     leading_dims = image.shape[:-3]
     C, H, W = image.shape[-3:]
     P = patch_size
+    if stride is None:
+        stride = P
 
     if (
         H % P != 0
@@ -104,14 +110,14 @@ def patchify(
         )
 
     N = (
-        int(H / P)
-        * int(W / P)
+        int((H - P + 1 + stride) // stride)
+        * int((W - P + 1 + stride) // stride)
     )
 
     patches = torch.nn.functional.unfold(
         input=image.reshape(-1, C, H, W),
         kernel_size=P,
-        stride=P,
+        stride=stride,
     )  # [prod(*), C*P*P, N]
     patches = torch.permute(patches, (0, 2, 1))  # [prod(*), N, C*P*P]