├── README.md ├── cae.py └── oc_nn.py /README.md: -------------------------------------------------------------------------------- 1 | # oc-nn-pytorch 2 | 3 | The repository consists of code for pytorch in [Anomaly Detection using One-Class Neural Networks](https://arxiv.org/abs/1802.06360) and includes only oc-nn-linear model. 4 | 5 | The code refers to [raghavchalapathy/oc-nn](https://github.com/raghavchalapathy/oc-nn). 6 | -------------------------------------------------------------------------------- /cae.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding: utf-8 -*- 3 | """ 4 | Created on Mon Jul 9 19:07:03 2018 5 | 6 | @author: seukgyo 7 | """ 8 | 9 | import os 10 | import torch 11 | import numpy as np 12 | import time 13 | import copy 14 | 15 | import torch.nn.functional as F 16 | 17 | from torchvision.datasets import MNIST 18 | from torch.utils.data import DataLoader 19 | from torch import nn 20 | 21 | 22 | """ 23 | DataLoader 24 | """ 25 | data_fold = 'data' 26 | 27 | if not os.path.isdir(data_fold): 28 | os.makedirs(data_fold) 29 | 30 | train_set = MNIST(root=data_fold, train=True, download=True) 31 | test_set = MNIST(root=data_fold, train=False, download=True) 32 | 33 | train_data = train_set.train_data.numpy() 34 | train_label = train_set.train_labels.numpy() 35 | 36 | normal_train = train_data[np.where(train_label==4), :, :] 37 | normal_train = normal_train.transpose(1, 0, 2, 3) 38 | 39 | normal_set = torch.FloatTensor(normal_train/255.) 40 | 41 | train_loader = DataLoader(normal_set, shuffle=True, batch_size=128, num_workers=4) 42 | 43 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 44 | 45 | dataset_size = len(normal_train) 46 | 47 | #%% 48 | """ 49 | Model 50 | """ 51 | cae_model_path = 'model/CAE.pth' 52 | 53 | class Encoder(nn.Module): 54 | def __init__(self): 55 | super(Encoder, self).__init__() 56 | 57 | self.conv1 = nn.Conv2d(1, 16, 3, stride=1, padding=1) 58 | self.conv2 = nn.Conv2d(16, 8, 3, stride=1, padding=1) 59 | self.conv3 = nn.Conv2d(8, 8, 3, stride=1, padding=1) 60 | 61 | self.pool1 = nn.MaxPool2d(2, stride=2) 62 | self.pool2 = nn.MaxPool2d(2, stride=2) 63 | 64 | self.dense1 = nn.Linear(392, 32) 65 | 66 | self.bn1 = nn.BatchNorm2d(16) 67 | self.bn2 = nn.BatchNorm2d(8) 68 | self.bn3 = nn.BatchNorm2d(8) 69 | 70 | def forward(self, img): 71 | x = self.conv1(img) 72 | x = self.bn1(x) 73 | x = F.elu(x) 74 | x = self.pool1(x) 75 | 76 | x = self.conv2(x) 77 | x = self.bn2(x) 78 | x = F.elu(x) 79 | x = self.pool2(x) 80 | 81 | x = self.conv3(x) 82 | x = self.bn3(x) 83 | x = F.elu(x) 84 | 85 | x = x.view(-1, 392) 86 | x = self.dense1(x) 87 | x = F.dropout(x, training=self.training) 88 | x = F.elu(x) 89 | 90 | return x 91 | 92 | class Decoder(nn.Module): 93 | def __init__(self): 94 | super(Decoder, self).__init__() 95 | 96 | self.deconv3 = nn.ConvTranspose2d(8, 8, 3, stride=1, padding=1) 97 | self.deconv2 = nn.ConvTranspose2d(8, 16, 3, stride=1, padding=1) 98 | self.deconv1 = nn.ConvTranspose2d(16, 1, 3, stride=1, padding=1) 99 | 100 | self.upsample2 = nn.Upsample(scale_factor=2, mode='bilinear') 101 | self.upsample1 = nn.Upsample(scale_factor=2, mode='bilinear') 102 | 103 | self.dense1 = nn.Linear(32, 392) 104 | 105 | self.bn3 = nn.BatchNorm2d(8) 106 | self.bn2 = nn.BatchNorm2d(16) 107 | 108 | def forward(self, encode): 109 | x = self.dense1(encode) 110 | x = F.dropout(x, training=self.training) 111 | x = F.elu(x) 112 | 113 | x = x.view(x.size(0), 8, 7, 7) 114 | 115 | x = self.deconv3(x) 116 | x = self.bn3(x) 117 | x = F.elu(x) 118 | 119 | x = self.upsample2(x) 120 | 121 | x = self.deconv2(x) 122 | x = self.bn2(x) 123 | x = F.elu(x) 124 | 125 | x = self.upsample1(x) 126 | 127 | x = self.deconv1(x) 128 | x = F.sigmoid(x) 129 | 130 | return x 131 | 132 | class CAE(nn.Module): 133 | def __init__(self): 134 | super(CAE, self).__init__() 135 | 136 | self.encoder = Encoder() 137 | self.decoder = Decoder() 138 | 139 | def forward(self, img): 140 | x = self.encoder(img) 141 | x = self.decoder(x) 142 | 143 | return x 144 | 145 | """ 146 | train 147 | """ 148 | if __name__ == '__main__': 149 | model = CAE() 150 | model = model.to(device) 151 | 152 | optimizer = torch.optim.Adam(model.parameters()) 153 | 154 | since = time.time() 155 | 156 | best_model_wts = copy.deepcopy(model.state_dict()) 157 | best_loss = 10000 158 | 159 | num_epochs = 100 160 | 161 | for epoch in range(num_epochs): 162 | print('Epoch {}/{}'.format(epoch, num_epochs - 1)) 163 | print('-' * 10) 164 | 165 | # Each epoch has a training and validation phase 166 | model.train() # Set model to training mode 167 | 168 | running_loss = 0.0 169 | 170 | # Iterate over data. 171 | for inputs in train_loader: 172 | inputs = inputs.to(device) 173 | 174 | # zero the parameter gradients 175 | optimizer.zero_grad() 176 | 177 | # forward 178 | # track history if only in train 179 | outputs = model(inputs) 180 | loss = F.mse_loss(inputs, outputs) 181 | 182 | # backward + optimize only if in training phase 183 | loss.backward() 184 | optimizer.step() 185 | 186 | # statistics 187 | running_loss += loss.item() * inputs.size(0) 188 | 189 | epoch_loss = running_loss / dataset_size 190 | 191 | print('Loss: {:.4f} '.format(epoch_loss)) 192 | 193 | # deep copy the model 194 | if epoch_loss < best_loss: 195 | best_loss = epoch_loss 196 | best_model_wts = copy.deepcopy(model.state_dict()) 197 | model.load_state_dict(best_model_wts) 198 | torch.save(model.state_dict(), cae_model_path) 199 | print() 200 | 201 | time_elapsed = time.time() - since 202 | print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60)) 203 | print('Best Loss: {:4f}'.format(best_loss)) 204 | 205 | -------------------------------------------------------------------------------- /oc_nn.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding: utf-8 -*- 3 | """ 4 | Created on Thu Jul 12 20:14:43 2018 5 | 6 | @author: seukgyo 7 | """ 8 | 9 | import os 10 | import numpy as np 11 | import time 12 | import copy 13 | 14 | import torch 15 | import torch.nn as nn 16 | import torch.nn.functional as F 17 | 18 | from torchvision.datasets import MNIST 19 | from torch import optim 20 | from torch.utils.data import DataLoader 21 | 22 | import cae 23 | 24 | from itertools import zip_longest 25 | import csv 26 | import matplotlib.pyplot as plt 27 | 28 | from sklearn.metrics import average_precision_score 29 | from sklearn.metrics import roc_auc_score 30 | from sklearn.metrics import precision_score 31 | 32 | #%% 33 | """ 34 | DataLoader 35 | """ 36 | data_fold = 'data' 37 | 38 | if not os.path.isdir(data_fold): 39 | os.makedirs(data_fold) 40 | 41 | train_set = MNIST(root=data_fold, train=True, download=True) 42 | test_set = MNIST(root=data_fold, train=False, download=True) 43 | 44 | train_data = train_set.train_data.numpy() 45 | train_label = train_set.train_labels.numpy() 46 | 47 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 48 | 49 | # normal class - 4 50 | class4 = train_data[np.where(train_label==4), :, :] 51 | class4 = class4.transpose(1, 0, 2, 3) 52 | 53 | rand_idx = np.random.choice(len(class4), 220) 54 | class4 = class4[rand_idx, :, :, :] 55 | 56 | # anomaly class - 0, 7, 9 57 | class0 = train_data[np.where(train_label==0), :, :] 58 | class0 = class0.transpose(1, 0, 2, 3) 59 | 60 | rand_idx = np.random.choice(len(class0), 5) 61 | class0 = class0[rand_idx, :, :, :] 62 | 63 | class7 = train_data[np.where(train_label==7), :, :] 64 | class7 = class7.transpose(1, 0, 2, 3) 65 | 66 | rand_idx = np.random.choice(len(class7), 3) 67 | class7 = class7[rand_idx, :, :, :] 68 | 69 | class9 = train_data[np.where(train_label==9), :, :] 70 | class9 = class9.transpose(1, 0, 2, 3) 71 | 72 | rand_idx = np.random.choice(len(class9), 3) 73 | class9 = class9[rand_idx, :, :, :] 74 | 75 | normal_class = class4 76 | anomaly_class = np.concatenate((class0, class7, class9), axis=0) 77 | 78 | """ 79 | pretrained model 80 | """ 81 | pretrained_model_path = 'model/CAE.pth' 82 | 83 | print('loading network...') 84 | model = cae.CAE() 85 | model.load_state_dict(torch.load(pretrained_model_path)) 86 | model = model.to(device) 87 | #%% 88 | 89 | model.eval() 90 | 91 | encoder = model.encoder 92 | 93 | """ 94 | forward encoder 95 | """ 96 | # normal encode 97 | normal_encode = [] 98 | 99 | for normal_img in normal_class: 100 | normal_img = np.reshape(normal_img, (1, 1, 28, 28)) 101 | normal_img = torch.FloatTensor(normal_img/255.) 102 | normal_img = normal_img.to(device) 103 | output = encoder(normal_img) 104 | 105 | output = output.cpu() 106 | output = output.detach().numpy() 107 | normal_encode.append(output) 108 | 109 | normal_encode = np.array(normal_encode) 110 | normal_encode = np.reshape(normal_encode, (normal_encode.shape[0], normal_encode.shape[2])) 111 | 112 | # anomaly encode 113 | anomaly_encode = [] 114 | 115 | for anomaly_img in anomaly_class: 116 | anomaly_img = np.reshape(anomaly_img, (1, 1, 28, 28)) 117 | anomaly_img = torch.FloatTensor(anomaly_img/255.) 118 | anomaly_img = anomaly_img.to(device) 119 | output = encoder(anomaly_img) 120 | 121 | output = output.cpu() 122 | output = output.detach().numpy() 123 | anomaly_encode.append(output) 124 | 125 | anomaly_encode = np.array(anomaly_encode) 126 | anomaly_encode = np.reshape(anomaly_encode, (anomaly_encode.shape[0], anomaly_encode.shape[2])) 127 | 128 | #%% 129 | """ 130 | train oc-nn 131 | """ 132 | """ 133 | oc-nn model 134 | """ 135 | oc_nn_model_path = 'model/oc_nn.pth' 136 | 137 | x_size = normal_encode.shape[1] 138 | h_size = 32 139 | y_size = 1 140 | 141 | class OC_NN(nn.Module): 142 | def __init__(self): 143 | super(OC_NN, self).__init__() 144 | 145 | self.dense_out1 = nn.Linear(x_size, h_size) 146 | self.out2 = nn.Linear(h_size, y_size) 147 | 148 | def forward(self, img): 149 | w1 = self.dense_out1(img) 150 | w2 = self.out2(w1) 151 | 152 | return w1, w2 153 | 154 | model = OC_NN() 155 | model.to(device) 156 | 157 | theta = np.random.normal(0, 1, h_size + h_size * x_size + 1) 158 | rvalue = np.random.normal(0, 1, (len(normal_encode), y_size)) 159 | nu = 0.04 160 | 161 | def nnscore(x, w, v): 162 | return torch.matmul(torch.matmul(x, w), v) 163 | 164 | def ocnn_loss(theta, x, nu, w1, w2, r): 165 | term1 = 0.5 * torch.sum(w1**2) 166 | term2 = 0.5 * torch.sum(w2**2) 167 | term3 = 1/nu * torch.mean(F.relu(r - nnscore(x, w1, w2))) 168 | term4 = -r 169 | 170 | return term1 + term2 + term3 + term4 171 | 172 | optimizer = optim.SGD(model.parameters(), lr=0.0001) 173 | 174 | dataset_size = len(normal_encode) 175 | normal_encode = torch.FloatTensor(normal_encode/255.) 176 | 177 | train_loader = DataLoader(normal_encode, batch_size=32, shuffle=True, num_workers=4, drop_last=True) 178 | 179 | since = time.time() 180 | 181 | best_model_wts = copy.deepcopy(model.state_dict()) 182 | best_loss = 10000 183 | 184 | num_epochs = 100 185 | 186 | 187 | for epoch in range(num_epochs): 188 | print('Epoch {}/{}'.format(epoch, num_epochs - 1)) 189 | print('-' * 10) 190 | 191 | # Each epoch has a training and validation phase 192 | model.train() # Set model to training mode 193 | 194 | running_loss = 0.0 195 | 196 | # Iterate over data. 197 | for inputs in train_loader: 198 | inputs = inputs.to(device) 199 | 200 | # zero the parameter gradients 201 | optimizer.zero_grad() 202 | 203 | # forward 204 | # track history if only in train 205 | w1, w2 = model(inputs) 206 | r = nnscore(inputs, w1, w2) 207 | loss = ocnn_loss(theta, inputs, nu, w1, w2, r) 208 | loss = loss.mean() 209 | 210 | # backward + optimize only if in training phase 211 | loss.backward() 212 | optimizer.step() 213 | 214 | # statistics 215 | running_loss += loss.item() * inputs.size(0) 216 | 217 | r = r.cpu().detach().numpy() 218 | r = np.percentile(r, q=100*nu) 219 | epoch_loss = running_loss / dataset_size 220 | 221 | print('Loss: {:.4f} '.format(epoch_loss)) 222 | print('Epoch = %d, r = %f'%(epoch+1, r)) 223 | 224 | # deep copy the model 225 | if epoch_loss < best_loss: 226 | best_loss = epoch_loss 227 | best_model_wts = copy.deepcopy(model.state_dict()) 228 | model.load_state_dict(best_model_wts) 229 | torch.save(model.state_dict(), oc_nn_model_path) 230 | 231 | 232 | print() 233 | 234 | time_elapsed = time.time() - since 235 | print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60)) 236 | print('Best Loss: {:4f}'.format(best_loss)) 237 | 238 | normal_encode = normal_encode.to(device) 239 | train_score = nnscore(normal_encode, w1, w2) 240 | train_score = train_score.cpu().detach().numpy() - r 241 | 242 | anomaly_encode = torch.FloatTensor(anomaly_encode) 243 | anomaly_encode = anomaly_encode.to(device) 244 | 245 | test_score = nnscore(anomaly_encode, w1, w2) 246 | test_score = test_score.cpu().detach().numpy() - r 247 | 248 | #%% 249 | """ 250 | Write Decision Scores to CSV 251 | """ 252 | decision_score_path = 'doc/oc-nn_linear.csv' 253 | 254 | print ('Writing file to ', decision_score_path) 255 | 256 | poslist = train_score.tolist() 257 | neglist = test_score.tolist() 258 | 259 | d = [poslist, neglist] 260 | export_data = zip_longest(*d, fillvalue='') 261 | with open(decision_score_path, 'w') as myfile: 262 | wr = csv.writer(myfile) 263 | wr.writerow(("Normal", "Anomaly")) 264 | wr.writerows(export_data) 265 | myfile.close() 266 | 267 | #%% 268 | """ 269 | Plot Decision Scores 270 | """ 271 | plt.plot() 272 | plt.title("One Class NN", fontsize="x-large", fontweight='bold'); 273 | plt.hist(train_score, bins=25, label='Normal') 274 | plt.hist(test_score, bins=25, label='Anomaly') 275 | 276 | #%% 277 | ## Obtain the Metrics AUPRC, AUROC, P@10 278 | y_train = np.ones(train_score.shape[0]) 279 | y_test = np.zeros(test_score.shape[0]) 280 | y_true = np.concatenate((y_train, y_test)) 281 | 282 | y_score = np.concatenate((train_score, test_score)) 283 | 284 | average_precision = average_precision_score(y_true, y_score) 285 | 286 | print('Average precision-recall score: {0:0.4f}'.format(average_precision)) 287 | 288 | roc_score = roc_auc_score(y_true, y_score) 289 | 290 | print('ROC score: {0:0.4f}'.format(roc_score)) 291 | 292 | def compute_precAtK(y_true, y_score, K = 10): 293 | 294 | if K is None: 295 | K = y_true.shape[0] 296 | 297 | # label top K largest predicted scores as + one's've 298 | 299 | idx = np.argsort(y_score) 300 | predLabel = np.zeros(y_true.shape) 301 | 302 | predLabel[idx[:K]] = 1 303 | 304 | prec = precision_score(y_true, predLabel) 305 | 306 | return prec 307 | 308 | prec_atk = compute_precAtK(y_true, y_score) 309 | 310 | print('Precision AtK: {0:0.4f}'.format(prec_atk)) 311 | --------------------------------------------------------------------------------