Spaces:
Running
Running
File size: 3,684 Bytes
f481275 |
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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import torch
import torch.nn as nn
import numpy as np
class SpectralFeatureExtractor(nn.Module):
def __init__(self, num_features):
super().__init__()
self.complex_weight = nn.Parameter(torch.randn(num_features, 2, dtype=torch.float32) * 0.02)
def forward(self, x):
B, L, C = x.shape
x = x.transpose(1, 2)
x_fft = torch.fft.rfft(x, dim=2, norm="ortho")
weight = torch.view_as_complex(self.complex_weight)
x_weighted = x_fft * weight.unsqueeze(0).unsqueeze(-1)
x_out = torch.fft.irfft(x_weighted, n=L, dim=2, norm="ortho")
return x_out.transpose(1, 2)
class ChannelMixing(nn.Module):
def __init__(self, num_features, reduction=4):
super().__init__()
self.fc1 = nn.Linear(num_features, num_features // reduction)
self.fc2 = nn.Linear(num_features // reduction, num_features)
self.act = nn.GELU()
def forward(self, x):
identity = x
x_pooled = x.mean(dim=1)
x_weighted = self.fc2(self.act(self.fc1(x_pooled)))
out = identity * x_weighted.unsqueeze(1)
return out
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=5000, dropout=0.2):
super().__init__()
self.dropout = nn.Dropout(p=dropout)
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0)
self.register_buffer("pe", pe)
def forward(self, x):
x = x + self.pe[:, : x.size(1), :]
return self.dropout(x)
class MineROINet(nn.Module):
def __init__(self, input_dim, d_model=64, nhead=2, num_layers=2, dim_feedforward=256, dropout=0.2, num_classes=3, seq_len=30):
super().__init__()
self.spectral = SpectralFeatureExtractor(input_dim)
self.channel_mix = ChannelMixing(input_dim)
self.input_projection = nn.Linear(input_dim, d_model) if input_dim != d_model else nn.Identity()
self.pos_encoder = PositionalEncoding(d_model, max_len=seq_len, dropout=dropout)
encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward,
dropout=dropout, activation="gelu", batch_first=True)
self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
self.classifier = nn.Sequential(
nn.LayerNorm(d_model),
nn.Dropout(dropout),
nn.Linear(d_model, d_model // 2),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(d_model // 2, num_classes),
)
def forward(self, seq):
seq = self.spectral(seq)
seq = self.channel_mix(seq)
seq = self.input_projection(seq)
seq = self.pos_encoder(seq)
z = self.transformer_encoder(seq)
pooled = z.mean(dim=1)
out = self.classifier(pooled)
return out
def create_model_30day(input_dim, num_classes=3):
return MineROINet(input_dim=input_dim, d_model=64, nhead=2, num_layers=2,
dim_feedforward=256, dropout=0.2, num_classes=num_classes, seq_len=30)
# def create_model_60day(input_dim, num_classes=3):
# return MineROINet(input_dim=input_dim, d_model=64, nhead=4, num_layers=2,
# dim_feedforward=256, dropout=0.2, num_classes=num_classes, seq_len=60) |