File size: 897 Bytes
8cee61b |
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 |
import torch
import torch.nn as nn
class Sadim54MStable(nn.Module):
def __init__(self, input_dim=5):
super().__init__()
# The exact architecture of SADIM 54M
self.encoder = nn.Sequential(
nn.Linear(input_dim, 1024), nn.ReLU(), nn.LayerNorm(1024),
nn.Linear(1024, 2048), nn.ReLU(), nn.LayerNorm(2048),
nn.Linear(2048, 512)
)
self.decoder = nn.Sequential(
nn.Linear(512, 2048), nn.ReLU(),
nn.Linear(2048, 1024), nn.ReLU(),
nn.Linear(1024, input_dim)
)
def forward(self, x):
return self.decoder(self.encoder(x))
def get_anomaly_score(self, x):
'''Calculate MSE loss for anomaly detection'''
with torch.no_grad():
reconstructed = self.forward(x)
error = torch.mean((x - reconstructed)**2, dim=1)
return error
|