Static dev

.github/workflows/black.yaml

@@ -0,0 +1,17 @@
+name: Python Black
+
+on: [push, pull_request]
+
+jobs:
+  lint:
+    name: Python Lint
+    runs-on: ubuntu-latest
+    steps:
+      - name: Setup Python
+        uses: actions/setup-python@v5
+      - name: Setup checkout
+        uses: actions/checkout@master
+      - name: Lint with Black
+        run: |
+          pip install black
+          black --diff --check src/embed_time

.github/workflows/tests.yaml

@@ -0,0 +1,28 @@
+name: Test
+
+on:
+  push:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10"]
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        pip install ".[dev]"
+    - name: Test with pytest
+      run: |
+        pytest --color=yes 
+    # Coverage should work out of the box for public repos. For private repos, more setup is likely required.
+    - name: Coverage
+      uses: codecov/codecov-action@v4

.gitignore

@@ -2,6 +2,13 @@
 *.zarr
 *.tiff
 *.tif
+**/data/
+**/mnist_data/
+**/data/
+**/mnist_data/
+/notebooks/splits
+/scripts/embed_time_runs
+/scripts/graphs
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -126,4 +133,8 @@ pyrepo
 .vscode/
 
 # OS Files
-.DS_Store
+.DS_Store
+
+#logs
+embed_time_runs/
+embed_time_static_runs/

README.md

@@ -5,6 +5,8 @@
 Create a new environment
 ```bash
 conda create -n embed_time python=3.10
-conda activate 3.10
+conda activate embed_time
 pip install -e .
+conda install -y pytorch-gpu cuda-toolkit=11.8 torchvision -c nvidia -c conda-forge -c pytorch
+
 ```

backup.py

@@ -0,0 +1,188 @@
+import torch
+from torch import nn, optim
+import torch.nn.functional as F
+import numpy as np
+
+class ResizeConv2d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, scale_factor, mode='nearest'):
+        super().__init__()
+        self.scale_factor = scale_factor
+        self.mode = mode
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=kernel_size//2)
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=self.scale_factor, mode=self.mode)
+        x = self.conv(x)
+        return x
+
+class ResizeArbitrary(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, out_size, mode='nearest'):
+        super().__init__()
+        self.out_size = out_size
+        self.mode = mode
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=kernel_size//2)
+
+    def forward(self, x):
+        x = F.interpolate(x, size=self.out_size, mode=self.mode)
+        x = torch.relu(self.conv(x))
+        return x
+
+class BasicBlockEnc(nn.Module):
+
+    def __init__(self, in_planes, stride=1):
+        super().__init__()
+
+        planes = in_planes*stride
+
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        if stride == 1:
+            self.shortcut = nn.Identity()
+        else:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes)
+            )
+
+    def forward(self, x):
+        out = torch.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out = out + self.shortcut(x)
+        out = torch.relu(out)
+        return out
+
+class BasicBlockDec(nn.Module):
+    def __init__(self, in_planes, stride=1):
+        super().__init__()
+        planes = int(in_planes/stride)
+
+        self.conv2 = nn.Conv2d(in_planes, in_planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(in_planes)
+        # self.bn1 could have been placed here, 
+        # but that messes up the order of the layers when printing the class
+
+        if stride == 1:
+            self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+            self.bn1 = nn.BatchNorm2d(planes)
+            self.shortcut = nn.Sequential()
+        else:
+            self.conv1 = ResizeConv2d(in_planes, planes, kernel_size=3, scale_factor=stride)
+            self.bn1 = nn.BatchNorm2d(planes)
+            self.shortcut = nn.Sequential(
+                ResizeConv2d(in_planes, planes, kernel_size=3, scale_factor=stride),
+                nn.BatchNorm2d(planes)
+            )
+    def forward(self, x):
+        out = torch.relu(self.bn2(self.conv2(x)))
+        out = self.bn1(self.conv1(out))
+        out = out + self.shortcut(x)
+        out = torch.relu(out)
+        return out
+
+
+class ResNet18Enc(nn.Module):
+
+    def __init__(self, num_Blocks=[2,2,2,2], z_dim=10, nc=3, linear_downsample_factor = 8):
+        super().__init__()
+        self.in_planes = 64
+        self.z_dim = z_dim
+        self.conv1 = nn.Conv2d(nc, 64, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.layer1 = self._make_layer(BasicBlockEnc, 64, num_Blocks[0], stride=1)
+        self.layer2 = self._make_layer(BasicBlockEnc, 128, num_Blocks[1], stride=2)
+        self.layer3 = self._make_layer(BasicBlockEnc, 256, num_Blocks[2], stride=2)
+        self.layer4 = self._make_layer(BasicBlockEnc, 512, num_Blocks[3], stride=2)
+        self.linear = nn.Linear(int(512*(128/2**len(num_Blocks))**2), 2 * z_dim, bias = False)
+
+    def _make_layer(self, BasicBlockEnc, planes, num_Blocks, stride):
+        strides = [stride] + [1]*(num_Blocks-1)
+        layers = []
+        for stride in strides:
+            layers += [BasicBlockEnc(self.in_planes, stride)]
+            self.in_planes = planes
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = torch.relu(self.bn1(self.conv1(x)))
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = torch.flatten(x, start_dim=1, end_dim=-1).unsqueeze(1)
+        x = torch.relu(self.linear(x))
+        mu, logvar = torch.chunk(x, 2, dim=2)
+        return mu, logvar
+
+class ResNet18Dec(nn.Module):
+
+    def __init__(self, spatial_dim_bottle, num_Blocks=[2,2,2,2], z_dim=10, nc=3, linear_downsample_factor =8):
+        super().__init__()
+        self.in_planes = 512
+        self.nc = nc
+        self.z_dim = z_dim
+
+
+        self.linear = nn.Linear(z_dim, 256)
+        self.firstconv = nn.Conv2d(1, 512, kernel_size=1)
+
+        self.layer4 = self._make_layer(BasicBlockDec, 256, num_Blocks[3], stride=2)
+        self.layer3 = self._make_layer(BasicBlockDec, 128, num_Blocks[2], stride=2)
+        self.layer2 = self._make_layer(BasicBlockDec, 64, num_Blocks[1], stride=2)
+        self.layer1 = self._make_layer(BasicBlockDec, 64, num_Blocks[0], stride=1)
+        self.conv1 = ResizeConv2d(64, nc, kernel_size=3, scale_factor=1)
+
+    def _make_layer(self, BasicBlockDec, planes, num_Blocks, stride):
+        strides = [stride] + [1]*(num_Blocks-1)
+        layers = []
+        for stride in reversed(strides):
+            layers += [BasicBlockDec(self.in_planes, stride)]
+        self.in_planes = planes
+        return nn.Sequential(*layers)
+
+    def forward(self, z):
+        x = torch.relu(self.linear(z))
+        x= x.view(-1, 1, 16,16)
+        x = self.firstconv(x)
+        x = self.layer4(x)
+        x = self.layer3(x)
+        x = self.layer2(x)
+        x = self.layer1(x)
+        x = torch.sigmoid(self.conv1(x))
+        return x
+
+
+class VAEResNet18_Linear(nn.Module):
+    def __init__(self, nc, z_dim, input_spatial_dim):
+        super().__init__()
+        self.in_spatial_shape = input_spatial_dim
+        self.spat_shape_bottle = self.compute_spatial_shape(4)
+        self.spat_shape_bottle = (self.spat_shape_bottle[0],self.spat_shape_bottle[1])
+        self.encoder = ResNet18Enc(nc=nc, z_dim=z_dim)
+        self.decoder = ResNet18Dec(nc=nc, z_dim=z_dim, spatial_dim_bottle=self.spat_shape_bottle)
+        self.enc_linear = nn.Sequential(
+
+        )
+
+    def forward(self, x):
+        mean, logvar = self.encoder(x)
+        z = self.reparameterize(mean, logvar)
+        x = self.decoder(z)
+        return x, z, mean, logvar
+
+    @staticmethod
+    def reparameterize(mean, logvar):
+        std = torch.exp(logvar / 2) # in log-space, squareroot is divide by two
+        epsilon = torch.randn_like(std)
+        return epsilon * std + mean
+
+    def compute_spatial_shape(self, level: int) -> tuple[int, int]:
+        # TODO Add warning when shape is odd before maxpool
+        spatial_shape = np.array(self.in_spatial_shape)
+        if level == 0:
+            return spatial_shape
+        spatial_shape = np.array(self.compute_spatial_shape(level-1)) // 2 
+        if any([s%2 != 0 for s in spatial_shape]):
+            raise ValueError("Can't Decode Because Input Dimension is Lost during Downsampling")
+        return spatial_shape