[論文理解] Unsupervised Anomaly Detection with Generative Adversarial Networks to Guide Marker Discovery

Unsupervised Anomaly Detection with Generative Adversarial Networks to Guide Marker Discovery

Intro

本文提出利用GAN來做異常檢測，大致思想為，先使用正常方式訓練生成器G，訓練完成后固定G的參數，給定一張圖片x，在latent space里去查詢與之最匹配的z，認為如果能找到一個合適的匹配，那么輸入x就應該服從訓練生成圖片的分布(近似認為是訓練GAN時正常圖片的分布)，即為正常圖片，否則為異常。基於此思想，作者設計了用於衡量匹配度的標准，用於區分正常樣本和異常和樣本。

Unsupervised Manifold Learning of Normal Anatomical Variability

這一部分就是正常的GAN的訓練過程，與之前其他異常檢測的工作一樣，這里采用分patch訓練的方式，將原圖划分為k個patch，通過對抗訓練將latent space采樣的z映射為image space的圖像。

常規的GAN訓練loss：

\[\mathop{min}_G \mathop{max}_D (D,G) = \mathbb{E}_{x \backsim p_{data}(x)}[logD(x)] + \mathbb{E_{z \backsim p_z(z)}}[log(1-D(G(z)))] \]

Mapping new Images to the Latent Space

GAN學習到的是如何將latent space中的采樣結果z映射到image space的圖像，但是反過的映射卻並沒有學習到，因此一個簡單的想法就是去查詢與輸入圖像最接近的對應z。

也即給定輸入x和一個已經經過訓練的生成器G，找到一個z，使得G(z)和x的差異最小。

寫成loss形式就是：

\[L_R(z_\gamma) = \sum |x - G(z_\gamma)| \]

顯然整個loss里只有一個參數z，可以通過對z求梯度進行迭代優化，不再贅述。上式在文章中稱之為Residual Loss，因為是殘差形式。

然而作者並沒有只使用這一個loss去優化z，因為上式只使用了GAN在訓練時的生成器，並沒將判別器加以利用，因此添加了判別器部分的discrimination loss，認為之前訓練的判別器也具備查詢能力，顯然如果判別器的輸出概率值在輸入圖片x和G(z)之間的差異較小，也可以認為這個z就是我們要找的z。文章覺得直接匹配輸出層的特征要更加合適，於是並不是直接使用的判別器的概率輸出，而是輸出特征。

\[L_G(z_\gamma) = \sum |f(x) - f(G(z_\gamma))| \]

這里f表示判別器特征提取函數。

因此，總的loss就可以寫為：

\[L(z_\gamma) = (1-\lambda) \cdot L_R(z_\gamma) + \lambda \cdot L_D(z_\gamma) \]

對z求梯度優化即可。

Detection of Anomalies

上面的loss可以作為評判樣本是否異常的標准，顯然loss越小，說明在訓練中見過相似的圖片，即輸入與訓練樣本同分布，所以判定為正常樣本，否則判定為異常樣本。

即將loss改寫為下式：

\[A(x) = (1-\lambda)\cdot R(x) + \lambda \cdot D(x) \]

通過殘差部分\(R(x) = |x - G(z_T)|\)可以得到異常區域，即差異的亮度可以表示異常區域和異常成都，而判別部分D(x)則可以反映置信水平。

效果見圖：

Coding

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
class G(nn.Module):
    def __init__(self):
        super(G,self).__init__()
        self.layer1 = nn.Sequential(
             nn.Linear(100,7*7*512),
             nn.BatchNorm1d(7*7*512),
             nn.ReLU(),
        )
        self.layer2 = nn.Sequential(
                        nn.ConvTranspose2d(512,256,3,2,1,1),
                        nn.BatchNorm2d(256),
                        nn.LeakyReLU(),
                        nn.ConvTranspose2d(256,128,3,1,1),
                        nn.BatchNorm2d(128),    
                        nn.LeakyReLU(),
            )
        self.layer3 = nn.Sequential(
                        nn.ConvTranspose2d(128,64,3,1,1),
                        nn.BatchNorm2d(64),    
                        nn.LeakyReLU(),
                        nn.ConvTranspose2d(64,1,3,2,1,1),
                        nn.Tanh()
            )
    def forward(self,z):
        out = self.layer1(z)
        out = out.view(out.size()[0],512,7,7)
        out = self.layer2(out)
        out = self.layer3(out)
        return out   

class D(nn.Module):
    def __init__(self):
        super(D,self).__init__()
        self.layer1 = nn.Sequential(
                        nn.Conv2d(1,8,3,padding=1),   # batch x 16 x 28 x 28
                        nn.BatchNorm2d(8),    
                        nn.LeakyReLU(),
                        nn.Conv2d(8,16,3,stride=2,padding=1),  # batch x 32 x 28 x 28
                        nn.BatchNorm2d(16),    
                        nn.LeakyReLU(),
                        #('max1',nn.MaxPool2d(2,2))   # batch x 32 x 14 x 14
        )
        self.layer2 = nn.Sequential(
                        nn.Conv2d(16,32,3,stride=2,padding=1),  # batch x 64 x 14 x 14
                        nn.BatchNorm2d(32),
                        nn.LeakyReLU(),
                        #nn.MaxPool2d(2,2),
                        nn.Conv2d(32,64,3,padding=1),  # batch x 128 x 7 x 7
                        nn.BatchNorm2d(64),
                        nn.LeakyReLU()
        )
        self.fc = nn.Sequential(
                        nn.Linear(64*7*7,1),
                        nn.Sigmoid()
        )
    def forward(self,x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.view(out.size()[0], -1)
        feature = out
        out = self.fc(out)
        return out,feature
class AnoGAN(nn.Module):
    def __init__(self,la = 0.1):
        super(AnoGAN,self).__init__()
        self.la = la
        self.G = G()
        self.D = D()
    def forward(self,x,z):
        fake_x = self.G(z)
        L_R = F.l1_loss(fake_x,x)
        _,f_x = self.D(x)
        _,f_fake_x = self.D(fake_x)
        L_D = F.l1_loss(f_fake_x,f_x)
        return (1-self.la) * L_R + self.la * L_D


if __name__ == "__main__":
    anogan = AnoGAN()
    #x = torch.randn(2,1,28,28)
    z = Variable(torch.nn.init.normal(torch.zeros(2,100),mean=0,std=0.1),requires_grad=True)
    #print(anogan(x,z))
    optimizer = torch.optim.Adam([z],lr=1e-3)
    datas = [torch.randn(2,1,28,28)]
    for epoch in range(2):
        for x in datas:
            optimizer.zero_grad()
            loss = anogan(x,z)
            print(loss)
            loss.backward()
            
            optimizer.step()