xianyupro / AnoGAN-pytorch

使用AnoGAN实现一个缺陷检测小demo

AnoGAN(Unsupervised Anomaly Detection with Generative Adversarial Networks to Guide Marker Discovery)

论文地址

下载链接

代码配套博客

博客地址

所需环境

• python
• pytorch
• Windows或Ubuntu或其它系统

文件结构

├── AnoGAN-pytorch: 根目录
  ├── data: mnist数据集下载保存地址  
  ├── result: 保存程序运行的结果
  ├── mnist_data.py: mnist数据集的读取
  ├── model.py: 生成器网络和判别器网络的搭建
  ├── train.py: 模型的训练与缺陷检测

 

训练步骤

1. 数据集下载

   点击☞☞☞ 下载mnist数据集至data文件夹中并解压到当前文件夹，data中文件目录如下：

   

   本次使用的mnist数据集为csv格式，方便读取。

2. 直接运行train.py文件。win直接右键运行即可，ubuntu使用python train.py执行。

 

注意事项

• 在mnist_data.py文件中定义mnist数据集存储路径，默认.\data\mnist_train.csv和.\data\mnist_test.csv。如若你的mnist数据集存放在别处，请自行修改。

• 在train.py文件中设置了--batch_size等等参数，这些参数只针对生成网络和对抗网络的训练，对于缺陷检测的测试，没有定义参数，大家可手动修改或自行添加相关参数。

  def parse_args():
      import argparse
      parser = argparse.ArgumentParser(description="pytorch fcn training")
      parser.add_argument("--device", default="cuda", help="training device")
      parser.add_argument("-b", "--batch_size", default=128, type=int)
      parser.add_argument("--epochs", default=300, type=int, metavar="N",
                          help="number of total epochs to train")
      args = parser.parse_args()
      return args

• 训练结果保存在result文件夹中，主要包括生成器生成的图片等等，具体大家可以阅读代码。

 

网络模型结构图

缺陷检测结果