File size: 7,573 Bytes
c76fc30 | 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 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 | """
model.py
========
LSTM-Autoencoder for REST API anomaly detection.
Supports two input modes:
- 'embedding' : token sequences (CSIC 2010) — integers → learned vectors
- 'numeric' : feature sequences (CIC-IDS2018, UNSW-NB15) — floats directly
The model trains only on normal traffic.
At inference, high reconstruction error = anomaly.
Author : K.A.D.S.D. Kandanaarachchi (2020/ICT/19)
Project: Detecting Anomalous REST API Traffic — IT4216
"""
import torch
import torch.nn as nn
class Encoder(nn.Module):
"""
Reads a session sequence and compresses it into a fixed-size context vector (the bottleneck).
Input : (batch, seq_len, input_size)
Output : (batch, hidden_size) — last hidden state only
"""
def __init__(
self,
input_size: int,
hidden_size: int,
num_layers: int = 2,
dropout: float = 0.2,
):
super().__init__()
self.lstm = nn.LSTM(
input_size=input_size,
hidden_size=hidden_size,
num_layers=num_layers,
batch_first=True,
dropout=dropout if num_layers > 1 else 0.0,
)
def forward(self, x):
_, (hidden, _) = self.lstm(x)
# hidden shape: (num_layers, batch, hidden_size)
# take only the top layer's hidden state
return hidden[-1]
class Decoder(nn.Module):
"""
Takes the bottleneck vector and reconstructs the original sequence step by step.
Input : (batch, hidden_size)
Output : (batch, seq_len, input_size)
"""
def __init__(
self,
hidden_size: int,
output_size: int,
seq_len: int,
num_layers: int = 2,
dropout: float = 0.2,
):
super().__init__()
self.seq_len = seq_len
self.lstm = nn.LSTM(
input_size=hidden_size,
hidden_size=hidden_size,
num_layers=num_layers,
batch_first=True,
dropout=dropout if num_layers > 1 else 0.0,
)
# Project LSTM output back to original feature size
self.output_layer = nn.Linear(hidden_size, output_size)
def forward(self, context):
# Repeat the context vector for each time step
# shape: (batch, seq_len, hidden_size)
repeated = context.unsqueeze(1).repeat(1, self.seq_len, 1)
out, _ = self.lstm(repeated)
# Project to output size at each timestep
return self.output_layer(out)
class LSTMAutoencoder(nn.Module):
"""
Full LSTM-Autoencoder for sequence anomaly detection.
Parameters
----------
input_mode : 'embedding' for token sequences (CSIC 2010)
'numeric' for float sequences (CIC-IDS2018, UNSW-NB15)
vocab_size : required when input_mode='embedding'
embed_dim : embedding dimension when input_mode='embedding'
input_size : number of features when input_mode='numeric'
hidden_size : LSTM hidden state size (bottleneck width)
seq_len : number of timesteps per session (window size)
num_layers : number of stacked LSTM layers
dropout : dropout between LSTM layers
"""
def __init__(
self,
input_mode: str = "numeric",
vocab_size: int = None,
embed_dim: int = 32,
input_size: int = 23,
hidden_size: int = 64,
seq_len: int = 5,
num_layers: int = 2,
dropout: float = 0.2,
):
super().__init__()
self.input_mode = input_mode
self.seq_len = seq_len
# Input layer depends on mode
if input_mode == "embedding":
assert vocab_size is not None, "vocab_size required for embedding mode"
self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
encoder_input_sz = embed_dim
decoder_output_sz = embed_dim
else:
self.embedding = None
encoder_input_sz = input_size
decoder_output_sz = input_size
self.encoder = Encoder(encoder_input_sz, hidden_size, num_layers, dropout)
self.decoder = Decoder(
hidden_size, decoder_output_sz, seq_len, num_layers, dropout
)
def forward(self, x):
"""
x shape:
embedding mode : (batch, seq_len) int64 token ids
numeric mode : (batch, seq_len, features) float32
"""
if self.input_mode == "embedding":
x = self.embedding(x) # → (batch, seq_len, embed_dim)
context = self.encoder(x)
reconstruction = self.decoder(context)
return reconstruction
def reconstruction_error(self, x):
"""
Compute per-sample mean squared error between input and reconstruction. Used as anomaly score.
Returns: (batch,) float tensor of error scores
"""
with torch.no_grad():
if self.input_mode == "embedding":
x_float = self.embedding(x).detach()
else:
x_float = x
recon = self.forward(x)
# MSE per sample: mean over seq_len and features
error = ((recon - x_float) ** 2).mean(dim=(1, 2))
return error
def build_model_csic2010(
vocab_size: int,
embed_dim: int = 32,
hidden_size: int = 64,
num_layers: int = 2,
seq_len: int = 5,
) -> LSTMAutoencoder:
"""Instantiate model configured for CSIC 2010 token sequences."""
return LSTMAutoencoder(
input_mode="embedding",
vocab_size=vocab_size,
embed_dim=embed_dim,
hidden_size=hidden_size,
seq_len=seq_len,
num_layers=num_layers,
)
def build_model_cicids2018(
n_features: int = 23,
hidden_size: int = 64,
num_layers: int = 2,
seq_len: int = 5,
) -> LSTMAutoencoder:
"""Instantiate model configured for CIC-IDS2018 flow features."""
return LSTMAutoencoder(
input_mode="numeric",
input_size=n_features,
hidden_size=hidden_size,
seq_len=seq_len,
num_layers=num_layers,
)
def build_model_unsw(
n_features: int = 20,
hidden_size: int = 64,
num_layers: int = 2,
seq_len: int = 5,
) -> LSTMAutoencoder:
"""Instantiate model configured for UNSW-NB15 flow features."""
return LSTMAutoencoder(
input_mode="numeric",
input_size=n_features,
hidden_size=hidden_size,
seq_len=seq_len,
num_layers=num_layers,
)
if __name__ == "__main__":
print("=== Smoke test — all three model configs ===\n")
# CSIC 2010
m1 = build_model_csic2010(vocab_size=51)
x1 = torch.randint(0, 51, (8, 5))
r1 = m1(x1)
e1 = m1.reconstruction_error(x1)
print(
f"CSIC 2010 input={tuple(x1.shape)} "
f"recon={tuple(r1.shape)} error={e1.mean():.4f}"
)
# CIC-IDS2018
m2 = build_model_cicids2018(n_features=23)
x2 = torch.randn(8, 5, 23)
r2 = m2(x2)
e2 = m2.reconstruction_error(x2)
print(
f"CIC-IDS2018 input={tuple(x2.shape)} "
f"recon={tuple(r2.shape)} error={e2.mean():.4f}"
)
# UNSW-NB15
m3 = build_model_unsw(n_features=20)
x3 = torch.randn(8, 5, 20)
r3 = m3(x3)
e3 = m3.reconstruction_error(x3)
print(
f"UNSW-NB15 input={tuple(x3.shape)} "
f"recon={tuple(r3.shape)} error={e3.mean():.4f}"
)
total = sum(p.numel() for p in m1.parameters())
print(f"\nCSIC 2010 model parameters: {total:,}")
total = sum(p.numel() for p in m2.parameters())
print(f"CIC-IDS2018 model parameters: {total:,}")
|