Upload modeling_neuroclr.py

modeling_neuroclr.py

@@ -0,0 +1,301 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import TransformerEncoder, TransformerEncoderLayer
+
+from transformers import PreTrainedModel
+from configuration_neuroclr import NeuroCLRConfig
+
+
+# --------------------------
+# SSL Encoder (per-ROI)
+# --------------------------
+class NeuroCLR(nn.Module):
+    def __init__(self, config: NeuroCLRConfig):
+        super().__init__()
+
+        encoder_layer = TransformerEncoderLayer(
+            d_model=config.TSlength,
+            dim_feedforward=2 * config.TSlength,
+            nhead=config.nhead,
+            batch_first=True,
+        )
+        self.transformer_encoder = TransformerEncoder(encoder_layer, config.nlayer)
+
+        self.projector = nn.Sequential(
+            nn.Linear(config.TSlength, config.projector_out1),
+            nn.BatchNorm1d(config.projector_out1),
+            nn.ReLU(),
+            nn.Linear(config.projector_out1, config.projector_out2),
+        )
+
+        self.normalize_input = config.normalize_input
+        self.pooling = config.pooling
+        self.TSlength = config.TSlength
+
+    def forward(self, x):
+        # x: [B, 1, 128]
+        if self.normalize_input:
+            x = F.normalize(x, dim=-1)
+
+        x = self.transformer_encoder(x)  # [B, 1, 128]
+
+        if self.pooling == "flatten":
+            h = x.reshape(x.shape[0], -1)  # [B, 128]
+        elif self.pooling == "mean":
+            h = x.mean(dim=1)
+        elif self.pooling == "last":
+            h = x[:, -1, :]
+        else:
+            raise ValueError(f"Unknown pooling='{self.pooling}'")
+
+        if h.shape[1] != self.TSlength:
+            raise ValueError(f"h dim {h.shape[1]} != TSlength {self.TSlength}")
+
+        z = self.projector(h)
+
+        return h, z
+
+
+# --------------------------
+# Your ResNet1D head (verbatim)
+# --------------------------
+class MyConv1dPadSame(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, groups=1):
+        super().__init__()
+        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, groups=groups)
+
+        self.kernel_size = kernel_size
+        self.stride = stride
+
+    def forward(self, x):
+        in_dim = x.shape[-1]
+        out_dim = (in_dim + self.stride - 1) // self.stride
+        p = max(0, (out_dim - 1) * self.stride + self.kernel_size - in_dim)
+        pad_left = p // 2
+        pad_right = p - pad_left
+        x = F.pad(x, (pad_left, pad_right), "constant", 0)
+        return self.conv(x)
+
+
+class MyMaxPool1dPadSame(nn.Module):
+    def __init__(self, kernel_size):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.stride = 1
+        self.max_pool = nn.MaxPool1d(kernel_size=kernel_size)
+
+    def forward(self, x):
+        in_dim = x.shape[-1]
+        out_dim = (in_dim + self.stride - 1) // self.stride
+        p = max(0, (out_dim - 1) * self.stride + self.kernel_size - in_dim)
+        pad_left = p // 2
+        pad_right = p - pad_left
+        x = F.pad(x, (pad_left, pad_right), "constant", 0)
+        return self.max_pool(x)
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, groups, downsample, use_bn, use_do, is_first_block=False):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.downsample = downsample
+        self.use_bn = use_bn
+        self.use_do = use_do
+        self.is_first_block = is_first_block
+
+        conv_stride = stride if downsample else 1
+
+        self.bn1 = nn.BatchNorm1d(in_channels)
+        self.relu1 = nn.ReLU()
+        self.do1 = nn.Dropout(p=0.75)
+        self.conv1 = MyConv1dPadSame(in_channels, out_channels, kernel_size, stride=conv_stride, groups=groups)
+
+        self.bn2 = nn.BatchNorm1d(out_channels)
+        self.relu2 = nn.ReLU()
+        self.do2 = nn.Dropout(p=0.75)
+        self.conv2 = MyConv1dPadSame(out_channels, out_channels, kernel_size, stride=1, groups=groups)
+
+        self.max_pool = MyMaxPool1dPadSame(kernel_size=conv_stride)
+
+    def forward(self, x):
+        identity = x
+
+        out = x
+        if not self.is_first_block:
+            if self.use_bn:
+                out = self.bn1(out)
+            out = self.relu1(out)
+            if self.use_do:
+                out = self.do1(out)
+        out = self.conv1(out)
+
+        if self.use_bn:
+            out = self.bn2(out)
+        out = self.relu2(out)
+        if self.use_do:
+            out = self.do2(out)
+        out = self.conv2(out)
+
+        if self.downsample:
+            identity = self.max_pool(identity)
+
+        if self.out_channels != self.in_channels:
+            identity = identity.transpose(-1, -2)
+            ch1 = (self.out_channels - self.in_channels) // 2
+            ch2 = self.out_channels - self.in_channels - ch1
+            identity = F.pad(identity, (ch1, ch2), "constant", 0)
+            identity = identity.transpose(-1, -2)
+
+        out += identity
+        return out
+
+
+class ResNet1D(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        base_filters,
+        kernel_size,
+        stride,
+        groups,
+        n_block,
+        n_classes,
+        downsample_gap=2,
+        increasefilter_gap=4,
+        use_bn=True,
+        use_do=True,
+        verbose=False
+    ):
+        super().__init__()
+        self.verbose = verbose
+        self.n_block = n_block
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.groups = groups
+        self.use_bn = use_bn
+        self.use_do = use_do
+        self.downsample_gap = downsample_gap
+        self.increasefilter_gap = increasefilter_gap
+
+        self.first_block_conv = MyConv1dPadSame(in_channels, base_filters, kernel_size=self.kernel_size, stride=1)
+        self.first_block_bn = nn.BatchNorm1d(base_filters)
+        self.first_block_relu = nn.ReLU()
+        out_channels = base_filters
+
+        self.basicblock_list = nn.ModuleList()
+        for i_block in range(self.n_block):
+            is_first_block = (i_block == 0)
+            downsample = (i_block % self.downsample_gap == 1)
+
+            if is_first_block:
+                in_ch = base_filters
+                out_ch = in_ch
+            else:
+                in_ch = int(base_filters * 2 ** ((i_block - 1) // self.increasefilter_gap))
+                if (i_block % self.increasefilter_gap == 0) and (i_block != 0):
+                    out_ch = in_ch * 2
+                else:
+                    out_ch = in_ch
+
+            block = BasicBlock(
+                in_channels=in_ch,
+                out_channels=out_ch,
+                kernel_size=self.kernel_size,
+                stride=self.stride,
+                groups=self.groups,
+                downsample=downsample,
+                use_bn=self.use_bn,
+                use_do=self.use_do,
+                is_first_block=is_first_block,
+            )
+            self.basicblock_list.append(block)
+            out_channels = out_ch
+
+        self.final_bn = nn.BatchNorm1d(out_channels)
+        self.final_relu = nn.ReLU(inplace=True)
+        self.dense = nn.Linear(out_channels, n_classes)
+
+    def forward(self, x):
+        out = self.first_block_conv(x)
+        if self.use_bn:
+            out = self.first_block_bn(out)
+        out = self.first_block_relu(out)
+
+        for block in self.basicblock_list:
+            out = block(out)
+
+        if self.use_bn:
+            out = self.final_bn(out)
+        out = self.final_relu(out)
+        out = out.mean(-1)
+        out = self.dense(out)
+        return out
+
+
+# --------------------------
+# HF model: encoder + ResNet1D head
+# --------------------------
+class NeuroCLRForSequenceClassification(PreTrainedModel):
+    """
+    Expected input x: [B, 200, 128]
+    - runs encoder per ROI: [B,1,128] -> h_r [B,128]
+    - stacks into H: [B,200,128]
+    - feeds ResNet1D: [B,200,128] -> logits
+    """
+    config_class = NeuroCLRConfig
+    base_model_prefix = "neuroclr"
+
+    def __init__(self, config: NeuroCLRConfig):
+        super().__init__(config)
+
+        self.encoder = NeuroCLR(config)
+
+        # Freeze the encoder
+        for p in self.encoder.parameters():
+            p.requires_grad = False
+
+        self.head = ResNet1D(
+            in_channels=config.n_rois,
+            base_filters=config.base_filters,
+            kernel_size=config.kernel_size,
+            stride=config.stride,
+            groups=config.groups,
+            n_block=config.n_block,
+            n_classes=config.num_labels,
+            downsample_gap=config.downsample_gap,
+            increasefilter_gap=config.increasefilter_gap,
+            use_bn=config.use_bn,
+            use_do=config.use_do,
+        )
+
+        self.post_init()
+
+    def forward(self, x: torch.Tensor, labels: torch.Tensor = None, **kwargs):
+        # x: [B, 200, 128]
+        if x.ndim != 3 or x.shape[1] != self.config.n_rois or x.shape[2] != self.config.TSlength:
+            raise ValueError(
+                f"Expected x shape [B,{self.config.n_rois},{self.config.TSlength}] but got {tuple(x.shape)}"
+            )
+
+        B, R, L = x.shape
+
+        # Encode each ROI independently (ROI-wise SSL)
+        hs = []
+        for r in range(R):
+            xr = x[:, r, :].unsqueeze(1)  # [B,1,128]
+            with torch.no_grad():
+                h, _ = self.encoder(xr)
+            # h, _ = self.encoder(xr)  # h: [B,128]
+            hs.append(h.unsqueeze(1))  # [B,1,128]
+
+        H = torch.cat(hs, dim=1)  # [B,200,128]
+
+        logits = self.head(H)     # head expects [B,200,128]
+        loss = None
+        if labels is not None:
+            loss = nn.CrossEntropyLoss()(logits, labels)
+
+        return {"loss": loss, "logits": logits}