Spaces:
Sleeping
Sleeping
| import torch | |
| from torch import nn | |
| import torch.nn.functional as F | |
| malware_classes = ['7ev3n', 'APosT', 'Adposhel', 'Agent', 'Agentb', 'Allaple', 'Alueron.gen!J', 'Amonetize', | |
| 'Androm', 'Bashlite', 'Bingoml', 'Blacksoul', 'BrowseFox', 'C2LOP.gen!g', 'Convagent', 'Copak', | |
| 'Delf', 'Dialplatform.B', 'Dinwod', 'Elex', 'Emotet', 'Escelar', 'Expiro', 'Fakerean', 'Fareit', | |
| 'Fasong', 'GandCrab', 'GlobelImposter', 'GootLoader', 'HLLP', 'HackKMS', 'Hlux', 'IcedId', 'Infy', | |
| 'Inject', 'Injector', 'InstallCore', 'KRBanker', 'Koadic', 'Kryptik', 'Kwampirs', 'Lamer', | |
| 'LemonDuck', 'Loki', 'Lolyda.AA1', 'Lolyda.AA2', 'Mimail', 'MultiPlug', 'Mydoom', 'Neoreklami', | |
| 'Neshta', 'NetWireRAT', 'Ngrbot', 'OnlinerSpambot', 'Orcus', 'Padodor', 'Plite', 'PolyRansom', | |
| 'QakBot', 'QtBot', 'Qukart', 'REvil', 'Ramdo', 'Regrun', 'Rekt Loader', 'Sakula', 'Salgorea', | |
| 'Scar', 'SelfDel', 'Small', 'Snarasite', 'Stantinko', 'Trickpak', 'Upantix', 'Upatre', 'VB', | |
| 'VBA', 'VBKrypt', 'VBNA', 'Vilsel', 'Vobfus', 'WBNA', 'Wecod', 'XTunnel', 'Zenpak', 'Zeus', 'benign'] | |
| class DenseLayer(nn.Module): | |
| def __init__(self, in_channels, growth_rate, bn_size): | |
| super(DenseLayer, self).__init__() | |
| self.bn1 = nn.BatchNorm2d(in_channels) | |
| self.conv1 = nn.Conv2d(in_channels, bn_size * growth_rate, kernel_size=1, bias=False) | |
| self.bn2 = nn.BatchNorm2d(bn_size * growth_rate) | |
| self.conv2 = nn.Conv2d(bn_size * growth_rate, growth_rate, kernel_size=3, padding=1, bias=False) | |
| def forward(self, x): | |
| out = self.conv1(F.relu(self.bn1(x))) | |
| out = self.conv2(F.relu(self.bn2(out))) | |
| return torch.cat([x, out], 1) | |
| class DenseBlock(nn.Module): | |
| def __init__(self, num_layers, in_channels, growth_rate, bn_size): | |
| super(DenseBlock, self).__init__() | |
| layers = [] | |
| for i in range(num_layers): | |
| layers.append(DenseLayer(in_channels + i * growth_rate, growth_rate, bn_size)) | |
| self.layers = nn.Sequential(*layers) | |
| def forward(self, x): | |
| return self.layers(x) | |
| class TransitionLayer(nn.Module): | |
| def __init__(self, in_channels, out_channels): | |
| super(TransitionLayer, self).__init__() | |
| self.bn = nn.BatchNorm2d(in_channels) | |
| self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False) | |
| self.pool = nn.AvgPool2d(kernel_size=2, stride=2) | |
| def forward(self, x): | |
| out = self.conv(F.relu(self.bn(x))) | |
| return self.pool(out) | |
| class MalwareNet(nn.Module): | |
| def __init__(self, growth_rate=32, block_config=(6, 12, 24, 16), num_init_features=64, bn_size=4, compression_rate=0.5, num_classes=87): | |
| super(MalwareNet, self).__init__() | |
| # First convolution | |
| self.features = nn.Sequential( | |
| nn.Conv2d(3, num_init_features, kernel_size=7, stride=2, padding=3, bias=False), | |
| nn.BatchNorm2d(num_init_features), | |
| nn.ReLU(inplace=True), | |
| nn.MaxPool2d(kernel_size=3, stride=2, padding=1) | |
| ) | |
| # Dense blocks | |
| num_features = num_init_features | |
| for i, num_layers in enumerate(block_config): | |
| block = DenseBlock(num_layers, num_features, growth_rate, bn_size) | |
| self.features.add_module(f'denseblock{i+1}', block) | |
| num_features += num_layers * growth_rate | |
| if i != len(block_config) - 1: | |
| transition = TransitionLayer(num_features, int(num_features * compression_rate)) | |
| self.features.add_module(f'transition{i+1}', transition) | |
| num_features = int(num_features * compression_rate) | |
| # Final batch norm | |
| self.features.add_module('norm5', nn.BatchNorm2d(num_features)) | |
| # Linear layer | |
| self.classifier = nn.Linear(num_features, num_classes) | |
| def forward(self, x): | |
| features = self.features(x) | |
| out = F.relu(features) | |
| out = F.adaptive_avg_pool2d(out, (1, 1)) | |
| out = torch.flatten(out, 1) | |
| out = self.classifier(out) | |
| return out |