| from torch import nn | |
| import torch.nn.functional as F | |
| class BadNet(nn.Module): | |
| """ Badnet model class based on the description of table1 of the paper with two convolution | |
| and two fully connected layers """ | |
| def __init__(self, input_size=3, output=10): | |
| super().__init__() | |
| self.input_size = input_size | |
| self.output = output | |
| self.conv1 = nn.Conv2d(in_channels=input_size, out_channels=16, kernel_size=(5, 5)) | |
| self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=(5, 5)) | |
| self.pool = nn.AvgPool2d(kernel_size=2, stride=2) | |
| if input_size == 3: | |
| self.fc_features = 800 | |
| else: | |
| self.fc_features = 512 | |
| self.fc1 = nn.Linear(self.fc_features, 512) | |
| self.fc2 = nn.Linear(512, output) | |
| def forward(self, x): | |
| x = self.conv1(x) | |
| x = F.relu(x) | |
| x = self.pool(x) | |
| x = self.conv2(x) | |
| x = F.relu(x) | |
| x = self.pool(x) | |
| x = x.contiguous().view(-1, self.fc_features) | |
| x = self.fc1(x) | |
| x = F.relu(x) | |
| x = self.fc2(x) | |
| x = F.softmax(x) | |
| return x | |