Initial release: Ultra-compact sub-millisecond gender classifier trained from scratch
9cd7b60 verified | import torch | |
| import torch.nn as nn | |
| class CompactGenderClassifier(nn.Module): | |
| """ | |
| Ultra-compact transformer binary gender classifier trained from scratch. | |
| Features: | |
| - Factorized Embeddings (vocab_size -> emb_dim -> hidden_dim) for ~50% embedding weight reduction | |
| - Pre-LN TransformerEncoder for fast, stable training dynamics | |
| - Masked Dual-Pooling: Concat(MeanPool, MaxPool) for subword affix sensitivity | |
| - Single-task binary classification head: MALE (0) vs FEMALE (1) | |
| - Zero BertPooler / token_type_ids bloat for clean ONNX execution | |
| """ | |
| def __init__( | |
| self, | |
| vocab_size=4000, | |
| emb_dim=64, | |
| hidden_dim=128, | |
| num_layers=2, | |
| num_heads=4, | |
| dim_feedforward=256, | |
| max_len=20, | |
| dropout=0.1 | |
| ): | |
| super().__init__() | |
| self.vocab_size = vocab_size | |
| self.emb_dim = emb_dim | |
| self.hidden_dim = hidden_dim | |
| self.max_len = max_len | |
| # Factorized Token Embeddings: cuts parameters from 4000x128=512k to 4000x64+64x128=264k | |
| self.token_emb = nn.Embedding(vocab_size, emb_dim, padding_idx=0) | |
| self.emb_proj = nn.Linear(emb_dim, hidden_dim, bias=False) | |
| self.pos_emb = nn.Embedding(max_len, hidden_dim) | |
| # Transformer Encoder with Pre-LN (norm_first=True) | |
| encoder_layer = nn.TransformerEncoderLayer( | |
| d_model=hidden_dim, | |
| nhead=num_heads, | |
| dim_feedforward=dim_feedforward, | |
| dropout=dropout, | |
| activation="gelu", | |
| batch_first=True, | |
| norm_first=True | |
| ) | |
| self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) | |
| self.norm = nn.LayerNorm(hidden_dim) | |
| self.dropout = nn.Dropout(dropout) | |
| # Dual Pooling: 128 (MeanPool) + 128 (MaxPool) = 256 | |
| self.classifier = nn.Linear(hidden_dim * 2, 2) | |
| self._init_weights() | |
| def _init_weights(self): | |
| """Standard truncated normal initialization for training from scratch.""" | |
| for m in self.modules(): | |
| if isinstance(m, nn.Linear): | |
| nn.init.trunc_normal_(m.weight, std=0.02) | |
| if m.bias is not None: | |
| nn.init.zeros_(m.bias) | |
| elif isinstance(m, nn.Embedding): | |
| nn.init.trunc_normal_(m.weight, std=0.02) | |
| if m.padding_idx is not None: | |
| with torch.no_grad(): | |
| m.weight[m.padding_idx].fill_(0) | |
| elif isinstance(m, nn.LayerNorm): | |
| nn.init.ones_(m.weight) | |
| nn.init.zeros_(m.bias) | |
| def forward(self, input_ids, attention_mask): | |
| # input_ids: [batch_size, seq_len] (int32 or int64) | |
| # attention_mask: [batch_size, seq_len] (int32 or int64, 1 for real tokens, 0 for pad) | |
| seq_len = input_ids.size(1) | |
| positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0) | |
| # Embeddings: token + positional | |
| x = self.emb_proj(self.token_emb(input_ids.long())) + self.pos_emb(positions) | |
| # PyTorch TransformerEncoder expects src_key_padding_mask=True for PAD tokens | |
| padding_mask = (attention_mask == 0) | |
| hidden = self.encoder(x, src_key_padding_mask=padding_mask) | |
| hidden = self.norm(hidden) # [batch_size, seq_len, hidden_dim] | |
| # 1. Masked Mean Pooling | |
| mask_expanded = attention_mask.unsqueeze(-1).float() | |
| sum_embeddings = torch.sum(hidden * mask_expanded, dim=1) | |
| sum_mask = torch.clamp(mask_expanded.sum(dim=1), min=1e-9) | |
| mean_pooled = sum_embeddings / sum_mask | |
| # 2. Masked Max Pooling (padded positions set to -1e9 so they are never max) | |
| hidden_masked = hidden.masked_fill(padding_mask.unsqueeze(-1), -1e9) | |
| max_pooled = torch.max(hidden_masked, dim=1).values | |
| # Concatenate features: 128 + 128 = 256 | |
| features = torch.cat([mean_pooled, max_pooled], dim=-1) | |
| features = self.dropout(features) | |
| logits = self.classifier(features) | |
| return logits | |