methods/classes/LandCodeTransformer.py - master - maratgaliulin/landcode_classifier

methods/classes/LandCodeTransformer.py
127 строк4 KB
maratgaliulin
first commit
26 янв 2026, 02:03
26 янв 2026, 02:0318d5031
import torch
import torch.nn as nn

from .PositionalEncoding import PositionalEncoding

class LandCodeTransformer(nn.Module):
    def __init__(self, tfidf_dim, num_classes, d_model=128, nhead=4, 
                 num_layers=3, dim_feedforward=256, dropout=0.3):
        """
        Трансформер - класс-ция земельных кодов
        
        Args:
            tfidf_dim: Размерность фичей
            num_classes: кол-во классов
        """
        super(LandCodeTransformer, self).__init__()
        
        self.tfidf_dim = tfidf_dim
        self.d_model = d_model
        
        
        self.tfidf_projection = nn.Linear(tfidf_dim, d_model)
        self.area_projection = nn.Linear(1, d_model)  
        
        
        self.pos_encoder = PositionalEncoding(d_model)
        
        
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
        
        self.classifier = nn.Sequential(
            nn.Linear(d_model * 3, 512),
            nn.LayerNorm(512),
            nn.ReLU(),
            nn.Dropout(dropout * 0.8),
            
            nn.Linear(512, 256),
            nn.LayerNorm(256),
            nn.ReLU(),
            nn.Dropout(dropout * 0.6),
            
            nn.Linear(256, 128),
            nn.LayerNorm(128),
            nn.ReLU(),
            nn.Dropout(dropout * 0.4),
            
            nn.Linear(128, num_classes)
        )
        
        # self.classifier = nn.Sequential(
        #     nn.Linear(tfidf_dim + 1, 256),  
        #     nn.BatchNorm1d(256),
        #     nn.ReLU(),
        #     nn.Dropout(0.5),
            
        #     nn.Linear(256, 128),
        #     nn.BatchNorm1d(128),
        #     nn.ReLU(),
        #     nn.Dropout(0.5),
            
        #     nn.Linear(128, num_classes)
        # )
        # ******** Dropout **********
        self.dropout = nn.Dropout(dropout)
        
        
        self._init_weights()
        
    def _init_weights(self):
        
        for name, param in self.named_parameters():
            if 'weight' in name and param.dim() > 1:
                nn.init.xavier_uniform_(param)
            elif 'bias' in name:
                nn.init.constant_(param, 0)
                
        
        nn.init.kaiming_uniform_(self.tfidf_projection.weight, nonlinearity='relu')
        nn.init.kaiming_uniform_(self.area_projection.weight, nonlinearity='relu')
        
    def forward(self, tfidf_features):
        """
        Продвиж.вперёд с предобработ.фичами
        
        Args:
            tfidf_features: [batch_size, tfidf_dim]
            area_features: [batch_size, 1]
        """
        batch_size = tfidf_features.size(0)
        
        
        tfidf_proj = self.tfidf_projection(tfidf_features)  
       
        
        
        sequence = torch.stack([tfidf_proj], dim=1)  
        
        
        sequence = self.pos_encoder(sequence)
        
        
        transformer_output = self.transformer(sequence)  
        
        
        mean_pooled = transformer_output.mean(dim=1)
        
        
        max_pooled, _ = transformer_output.max(dim=1)
        
        
        weights = torch.softmax(transformer_output.mean(dim=-1), dim=-1)  
        weighted = (transformer_output * weights.unsqueeze(-1)).sum(dim=1)
        
        
        combined = torch.cat([mean_pooled, max_pooled, weighted], dim=1)  
        
        
        logits = self.classifier(combined)
        
        return logits