| import torch | |
| import torch.nn as nn | |
| class CyberClassifier(nn.Module): | |
| """Binary classification model for cybersecurity detection.""" | |
| def __init__(self, input_dim): | |
| super(CyberClassifier, self).__init__() | |
| self.normalizer = nn.LayerNorm(input_dim) | |
| self.fc1 = nn.Linear(input_dim, 512) | |
| self.bn1 = nn.BatchNorm1d(512) | |
| self.fc2 = nn.Linear(512, 256) | |
| self.bn2 = nn.BatchNorm1d(256) | |
| self.fc3 = nn.Linear(256, 128) | |
| self.bn3 = nn.BatchNorm1d(128) | |
| self.fc4 = nn.Linear(128, 1) | |
| def forward(self, x): | |
| """Forward pass through the network.""" | |
| x = self.normalizer(x) | |
| x = self.fc1(x) | |
| x = self.bn1(x) | |
| x = torch.relu(x) | |
| x = self.fc2(x) | |
| x = self.bn2(x) | |
| x = torch.relu(x) | |
| x = self.fc3(x) | |
| x = self.bn3(x) | |
| x = torch.relu(x) | |
| x = self.fc4(x) | |
| x = torch.sigmoid(x) | |
| return x | |
| def load_model(checkpoint_path, input_dim, device='cpu'): | |
| """ | |
| Load the model from checkpoint. | |
| Args: | |
| checkpoint_path: Path to model checkpoint (.pt file) | |
| input_dim: Input dimension (e.g., 768) | |
| device: 'cuda' or 'cpu' | |
| Returns: | |
| model: Loaded model ready for inference | |
| """ | |
| model = CyberClassifier(input_dim).to(device) | |
| checkpoint = torch.load(checkpoint_path, map_location=device) | |
| model.load_state_dict(checkpoint['model_state_dict']) | |
| model.eval() | |
| return model | |