File size: 2,507 Bytes
24fe966 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | import torch
import torch.nn as nn
from transformers import PreTrainedModel
from configuration_chronogrid import ChronoGridConfig
def make_cnn_backbone(img_size):
return nn.Sequential(
# Block 1 — kernel-3 (captures short transients)
nn.Conv1d(img_size, 64, kernel_size=3, padding=1),
nn.BatchNorm1d(64), nn.ReLU(),
# Block 2 — kernel-5 (captures broader waveform distortions)
nn.Conv1d(64, 128, kernel_size=5, padding=2),
nn.BatchNorm1d(128), nn.ReLU(), nn.MaxPool1d(2), # → [B, 128, 113]
# Block 3 — kernel-3 (deep feature abstraction)
nn.Conv1d(128, 256, kernel_size=3, padding=1),
nn.BatchNorm1d(256), nn.ReLU(), nn.MaxPool1d(2), # → [B, 256, 56]
)
class _ScaledDotAttn(nn.Module):
'''Scaled dot-product multi-head self-attention with residual + LayerNorm.'''
def __init__(self, d_model=256, num_heads=8, dropout=0.1):
super().__init__()
self.mha = nn.MultiheadAttention(d_model, num_heads, dropout=dropout,
batch_first=True)
self.norm = nn.LayerNorm(d_model)
def forward(self, x):
out, w = self.mha(x, x, x, need_weights=True, average_attn_weights=True)
return self.norm(x + out), w # w: [B, seq, seq]
class ChronoGridModelForSequenceClassification(PreTrainedModel):
config_class = ChronoGridConfig
def __init__(self, config):
super().__init__(config)
self.num_classes = config.num_classes
# 1. Initialize your architecture
self.cnn = make_cnn_backbone(config.img_size)
self.bilstm = nn.LSTM(256, 128, num_layers=2, batch_first=True, bidirectional=True, dropout=0.3)
self.attention = _ScaledDotAttn(d_model=256, num_heads=8)
self.head = nn.Sequential(
nn.Dropout(0.3),
nn.Linear(256, 128), nn.ReLU(),
nn.Linear(128, self.num_classes),
)
def forward(self, x, labels=None):
# x shape: [B, 227, 227]
x = self.cnn(x).permute(0, 2, 1)
x, _ = self.bilstm(x)
x, attn_w = self.attention(x)
feat = x.mean(dim=1)
logits = self.head(feat)
loss = None
if labels is not None:
loss_fct = nn.CrossEntropyLoss()
loss = loss_fct(logits.view(-1, self.num_classes), labels.view(-1))
return (loss, logits) if loss is not None else (logits,)
|