Add PyTorch original: model_state_dict.pt + model_def.py + torch_source.json

README.md

@@ -38,9 +38,9 @@ Endoplasmic reticulum (er) from FIB-SEM at 8 nm, 3D residual U-Net (InstanceNorm
 
 | File | Format | Usage |
 |---|---|---|
-| `model.pt2` | torch.export | `torch.export.load("model.pt2")` (trainable via `CellmapModel.train()`) |
+| `model.pt` | PyTorch pickle | `torch.load("model.pt")` |
 | `model.ts` | TorchScript | `torch.jit.load("model.ts")` |
-| `model.onnx` (+ `model.onnx.data`) | ONNX | `onnxruntime.InferenceSession("model.onnx")` |
+| `model.onnx` | ONNX | `onnxruntime.InferenceSession("model.onnx")` |
 | `metadata.json` | JSON | Model metadata |
 
 ## Usage
@@ -67,7 +67,7 @@ Or download from this repo and load directly:
 from huggingface_hub import snapshot_download
 from cellmap_models.model_export.cellmap_model import CellmapModel
 
-path = snapshot_download(repo_id="cellmap/er-resunet-setup-67")
+path = snapshot_download(repo_id="er_resunet_setup_67")
 model = CellmapModel(path)
 ```
 
@@ -79,3 +79,18 @@ Marwan Zouinkhi
 
 - [cellmap-models](https://github.com/janelia-cellmap/cellmap-models)
 - [CellMap Project](https://www.janelia.org/project-team/cellmap)
+
+## PyTorch original (for finetuning / re-export)
+
+`model_state_dict.pt` (plain tensors, `weights_only=True`) + the module code (`model_def.py`); `torch_source.json`
+names the class and constructor arguments.
+
+```python
+import json, importlib.util, torch
+src = json.load(open("torch_source.json"))
+spec = importlib.util.spec_from_file_location("model_def", src["module"]); mod = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(mod)
+model = getattr(mod, src["class"])(**src["kwargs"])
+model.load_state_dict(torch.load("model_state_dict.pt", weights_only=True))
+model.eval()
+```

model_def.py

@@ -0,0 +1,118 @@
+"""Custom 3D residual U-Net (~10.7M params), valid-padding throughout.
+
+Valid (unpadded) convolutions are used instead of same-padding because this setup is intended to run
+under blockwise inference (daisy/gunpowder tiling adjacent crops across a large volume). Same-padding
+convs synthesize zero-context at every block's border, which produces seam artifacts at block
+boundaries; valid convs give every output voxel genuine surrounding context, matching the convention
+already used by fly_organelles.model.StandardUnet and fish_trainer.model.Unet3d.
+
+Verified (base_channels=40, factor=2, blocks_per_level=1): 10,702,493 params,
+input (148,148,148) -> output (60,60,60).
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def center_crop_to(x, target):
+    diffs = [x.shape[i] - target.shape[i] for i in (2, 3, 4)]
+    starts = [d // 2 for d in diffs]
+    return x[
+        :,
+        :,
+        starts[0] : starts[0] + target.shape[2],
+        starts[1] : starts[1] + target.shape[3],
+        starts[2] : starts[2] + target.shape[4],
+    ]
+
+
+class ResBlock(nn.Module):
+    """Two valid 3x3x3 convs + InstanceNorm + ReLU, with a crop-and-add identity skip
+    (1x1x1 conv on the skip path when channel count changes; identity is center-cropped
+    to match the main path's spatially-shrunk output before the residual add)."""
+
+    def __init__(self, c_in, c_out):
+        super().__init__()
+        self.conv1 = nn.Conv3d(c_in, c_out, 3, padding=0)
+        self.norm1 = nn.InstanceNorm3d(c_out, affine=True)
+        self.conv2 = nn.Conv3d(c_out, c_out, 3, padding=0)
+        self.norm2 = nn.InstanceNorm3d(c_out, affine=True)
+        self.act = nn.ReLU(inplace=True)
+        self.skip = nn.Conv3d(c_in, c_out, 1) if c_in != c_out else nn.Identity()
+
+    def forward(self, x):
+        idn = self.skip(x)
+        h = self.act(self.norm1(self.conv1(x)))
+        h = self.norm2(self.conv2(h))
+        idn = center_crop_to(idn, h)
+        return self.act(h + idn)
+
+
+class Down(nn.Module):
+    """Exact 2x downsample (kernel==stride, divides evenly, no cropping needed)."""
+
+    def __init__(self, c_in, c_out):
+        super().__init__()
+        self.conv = nn.Conv3d(c_in, c_out, 2, stride=2)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class Up(nn.Module):
+    """Non-learnable nearest-neighbor upsample (exact 2x) + 1x1x1 channel-reduce conv.
+    Deliberately non-learnable to match this codebase's own constant_upsample=True choice in
+    StandardUnet (avoids checkerboard/seam artifacts from learned transposed convs)."""
+
+    def __init__(self, c_in, c_out):
+        super().__init__()
+        self.reduce = nn.Conv3d(c_in, c_out, 1)
+
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=2, mode="nearest")
+        return self.reduce(x)
+
+
+class ResUNet3DValid(nn.Module):
+    def __init__(self, out_channels, base_channels=40, factor=2, blocks_per_level=1, in_channels=1):
+        super().__init__()
+        C = [base_channels * (factor**i) for i in range(4)]  # 3 down/up levels + bottleneck
+        self.C = C
+
+        self.stem = ResBlock(in_channels, C[0])
+
+        self.downs = nn.ModuleList()
+        self.enc_blocks = nn.ModuleList()
+        for i in range(3):
+            self.downs.append(Down(C[i], C[i + 1]))
+            self.enc_blocks.append(nn.ModuleList([ResBlock(C[i + 1], C[i + 1]) for _ in range(blocks_per_level)]))
+
+        self.ups = nn.ModuleList()
+        self.dec_blocks = nn.ModuleList()
+        for i in reversed(range(3)):
+            self.ups.append(Up(C[i + 1], C[i]))
+            blocks = [ResBlock(C[i] * 2, C[i])] + [ResBlock(C[i], C[i]) for _ in range(blocks_per_level - 1)]
+            self.dec_blocks.append(nn.ModuleList(blocks))
+
+        self.final_conv = nn.Conv3d(C[0], out_channels, 1)
+
+    def forward(self, input):
+        skips = []
+        x = self.stem(input)
+        skips.append(x)
+        for down, enc in zip(self.downs, self.enc_blocks):
+            x = down(x)
+            for blk in enc:
+                x = blk(x)
+            skips.append(x)
+        skips.pop()  # bottleneck itself is not a skip target
+
+        for up, dec, skip in zip(self.ups, self.dec_blocks, reversed(skips)):
+            x = up(x)
+            skip_c = center_crop_to(skip, x)
+            x = torch.cat([x, skip_c], dim=1)
+            for blk in dec:
+                x = blk(x)
+
+        return self.final_conv(x)

model_state_dict.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0905b254f38ef5d3ae46335b92bf96025695162d3c7b909a2a5842f2826d9c57
+size 42832173

torch_source.json

@@ -0,0 +1,12 @@
+{
+ "module": "model_def.py",
+ "class": "ResUNet3DValid",
+ "kwargs": {
+  "out_channels": 1,
+  "base_channels": 40,
+  "factor": 2,
+  "blocks_per_level": 1
+ },
+ "state_dict": "model_state_dict.pt",
+ "torch_version": "2.9.1+cu128"
+}