一、前言

本项目为百度论文复现赛《FastFlow: Unsupervised Anomaly Detection and Localization via 2D Normalizing Flows》论文复现代码。
github地址:fastflow_paddle
参考repo:anomalib

依赖环境:

  • paddlepaddle-gpu2.3.2
  • python3.7

在MVTec-AD数据集下,复现精度:
预训练权重下载:logs (提取码:3ra1)下载后,放置到work/results文件夹下即可

FastFlow(ResNet18 )image-level AUCpixel-level AUC
论文97.997.2
复现98.097.2

二、模型背景及其介绍

论文中提出了一种图像缺陷异常检测模型,可以不依赖于异常数据来检测未知的异常缺陷。具体来说,提出了使用2D normalizing flow的FastFLow,并使用它来估计概率分布。FastFlow可以作为plug-in模块,与任意的深度特征提取器(如ResNet和Vision Transformer)一起使用,用于无监督异常检测和定位。在训练阶段,FastFlow学习将输入的视觉特征转化为可处理的分布,并在测试阶段得到异常的似然(即概率)

在这里插入图片描述

FastFlow包含两个模块,输入normal images,首先使用预训练的Feature Extractor(如Resnet18)提取特征,然后使用FastFlow学习特征分布。

三、数据集

MVTec AD是MVtec公司提出的一个用于异常检测的数据集。与之前的异常检测数据集不同,该数据集模仿了工业实际生产场景,并且主要用于unsupervised anomaly detection。数据集为异常区域都提供了像素级标注,是一个全面的、包含多种物体、多种异常的数据集。数据集包含不同领域中的五种纹理以及十种物体,且训练集中只包含正常样本,测试集中包含正常样本与缺陷样本,因此需要使用无监督方法学习正常样本的特征表示,并用其检测缺陷样本。

数据集下载链接:AiStudio数据集

四、运行

1、解压预训练数据

# 解压数据集
!tar xvf /home/aistudio/data/data116034/mvtec_anomaly_detection.tar.xz -C /home/aistudio/work/data/

#   - data目录结构为:
#     ```
#     data
#         └── bottle
#             ├── ground_truth
#             ├── test
#             ├── train
#         └── cable
#             ├── ground_truth
#             ├── test
#             ├── train 
#         └── ...
#     ```

2、安装环境包

%cd /home/aistudio/work
%pip install --upgrade pip
%pip install -r requirements.txt

3、训练

# 单类别全量数据训练
%cd /home/aistudio/work
!python train.py  -cfg ./configs/resnet18.yaml --data ./data --exp_dir exp -cat bottle
# 所有类别全量数据训练
%cd /home/aistudio/work
!sh train.sh
# 少量数据训练
%cd /home/aistudio/work
!python train.py  -cfg ./configs/resnet18.yaml --data ./lite_data --exp_dir exp -cat bottle

*模型搭建核心代码展示

class FastFlow(nn.Layer):
    '''
    FastFLow model
    '''
    def __init__(self,
        flow_steps = 8,
        input_size = 256,
        conv3x3_only=True,
        hidden_ratio=1.0,
        use_norm = True,
        momentum = 0.95,
        channels = [64, 128, 256],
        scales = [4, 8, 16],
        clamp = 2.0,
        ):
        super().__init__()
        #### Moudle1: Encoder - resnet18
        self.feature_extractor = resnet18(pretrained=True)          
        for param in self.feature_extractor.parameters():
            param.stop_gradient = True

        self.input_size = input_size
        #### Moudle2: Norm - BatchNorm
        self.Norm = use_norm
        if self.Norm:
            self.norms = nn.LayerList()
            for in_channels, scale in zip(channels, scales):                
                self.norms.append(
                    nn.BatchNorm2D(
                        in_channels, momentum=momentum
                    )
                )
        #### Moudle3: 2D Normalizing Flows - fastflow
        self.nf_flows = nn.LayerList()
        for in_channels, scale in zip(channels, scales):
            self.nf_flows.append(
                nf_fast_flow(
                    [in_channels, int(input_size / scale), int(input_size / scale)],
                    conv3x3_only=conv3x3_only,
                    hidden_ratio=hidden_ratio,
                    flow_steps=flow_steps,
                    clamp=clamp,
                )
            )

    def forward(self, x):
        
        ##step 1: encode feature
        self.feature_extractor.eval()
        features = self.feature_extractor(x)
        
        ## step 2: norm features
        if self.Norm:
            features = [self.norms[i](feature) for i, feature in enumerate(features)]
        
        
        loss = 0
        outputs = []
        ## step3: fastflow
        for i, feature in enumerate(features):
            output, log_jac_dets = self.nf_flows[i](feature)
            ### loss = -logP(y) = -(logP(z) + jac)
            loss +=  (paddle.mean(
                            0.5 * paddle.sum(output**2, axis=(1, 2, 3)) -log_jac_dets
                        ) / (output.shape[1] * output.shape[2] * output.shape[3]))

            
            outputs.append(output)
        ### step4 :post process
        ret = {"loss": loss}
        if not self.training:
            anomaly_map_list = []
            for output in outputs:
                log_prob = -paddle.mean(output**2, axis=1, keepdim=True) * 0.5  ###logP(z)
                prob = paddle.exp(log_prob) ###P(z)
                ### get the final probability map and upsample it to the input image resolution using bilinear interpolation.
                a_map = F.interpolate(
                    1-prob,
                    size=[self.input_size, self.input_size],
                    mode="bilinear",
                    align_corners=True,
                )

                anomaly_map_list.append(a_map)

            anomaly_map_list = paddle.stack(anomaly_map_list, axis=-1)
            anomaly_map = paddle.mean(anomaly_map_list, axis=-1)

            ret["anomaly_map"] = anomaly_map
        return ret

* 训练核心代码展示

def train_one_epoch(dataloader, model, optimizer, epoch, category, summary_writer, mylogger):
    '''
    train 1 epoch
    '''
    model.train()
    loss_meter = AverageMeter()
    pbar = tqdm(dataloader)

    train_reader_cost = 0.0 ## cost time on loading train dataset
    train_run_cost = 0.0
    total_samples = 0
    reader_start = time.time()
    for step, (data, file) in enumerate(pbar):
        # forward
        train_reader_cost += time.time() - reader_start
        train_start = time.time()
        ret = model(data)
        loss = ret["loss"]
        summary_writer.add_scalar("train_loss/iter", loss.item(), step)
        # backward  
        loss.backward()
        # optimizer
        optimizer.step()
        optimizer.clear_grad()
        
        train_run_cost += time.time() - train_start
        total_samples += data.shape[0]
        # log
        loss_meter.update(loss.item())
        message = "{}  Epoch {} - iter:{}/{} - lr:{} - loss = {:.3f}/{:.3f} (last/avg) - avg_reader_cost: {:.5f} sec - avg_batch_cost: {:.5f} sec - avg_samples: {} - avg_ips: {:.5f} images/sec".format(
              category, epoch + 1, step + 1, len(pbar), optimizer.get_lr(), loss_meter.val, loss_meter.avg, train_reader_cost, train_reader_cost + train_run_cost, total_samples, total_samples / (train_reader_cost + train_run_cost)
            )
        train_reader_cost = 0.0
        train_run_cost = 0.0
        total_samples = 0
        mylogger.print(message)
        pbar.set_postfix_str(message)
        reader_start = time.time()

4、评估

# 全量数据模型评估
%cd /home/aistudio/work
!python eval.py  -cfg ./configs/resnet18.yaml --data ./data -cat all --exp_dir exp
# 少量数据模型评估
%cd /home/aistudio/work
!python eval.py  -cfg ./configs/resnet18.yaml --data ./lite_data -cat bottle --exp_dir exp/
/home/aistudio/work
/opt/conda/envs/python35-paddle120-env/lib/python3.7/site-packages/setuptools/depends.py:2: DeprecationWarning: the imp module is deprecated in favour of importlib; see the module's documentation for alternative uses
  import imp
Model A.D. Param#: Tensor(shape=[1], dtype=int64, place=Place(cpu), stop_gradient=False,
       [4896728])
EVAL:  image-auc : 1.000000 pixel-auc:0.994000
bottle           pixel auroc:0.994   image auroc:1.000
all              pixel auroc:0.994   image auroc:1.000

* 评估核心代码展示

def eval_once(dataloader, model, eval = True):
    '''
    function:eval on test dataset
    
    return:
        auroc_px: pixel_level auc
        auroc_sp: image_level auc
    '''
    model.eval()
    gt_list_px = []
    pr_list_px = []
    gt_list_sp = []
    pr_list_sp = []
    preds = []
    gts = []

    for data, targets in dataloader:
        targets = targets.cpu().numpy().astype(int)

        with paddle.no_grad():
            ret = model(data)

        outputs = ret["anomaly_map"].detach().cpu()
        outputs = outputs.numpy()

        preds.append(outputs)
        gts.append(targets)
    
    preds = np.concatenate(preds, axis = 0)
    gts = np.concatenate(gts, axis = 0)

    targets = gts

    outputs = preds
    for i in range(targets.shape[0]):
        gt_list_sp.append(np.max(targets[i]))
        pr_list_sp.append(np.max(outputs[i]))
        if eval:
            outputs[i] = gaussian_filter(outputs[i], sigma=6)

        gt_list_px.extend(targets[i].ravel())
        pr_list_px.extend(outputs[i].ravel())
        
    auroc_px = round(roc_auc_score(gt_list_px, pr_list_px), 3)
    auroc_sp = round(roc_auc_score(gt_list_sp, pr_list_sp), 3)

    print("EVAL:  image-auc : {:.6f} pixel-auc:{:.6f}".format(auroc_sp, auroc_px))
    
    return auroc_px, auroc_sp

5、预测

# 模型预测
%cd /home/aistudio/work
!python predict.py -cfg ./configs/resnet18.yaml --category bottle --image_path images/bottle_good.png --exp_dir exp
/home/aistudio/work
/opt/conda/envs/python35-paddle120-env/lib/python3.7/site-packages/setuptools/depends.py:2: DeprecationWarning: the imp module is deprecated in favour of importlib; see the module's documentation for alternative uses
  import imp
Model A.D. Param#: Tensor(shape=[1], dtype=int64, place=Place(cpu), stop_gradient=False,
       [4896728])
Normal - score:  0.275
# 基于推理引擎导出模型
%cd /home/aistudio/work
!python deploy/export_model.py --exp_dir exp --category bottle -cfg configs/resnet18.yaml --save_inference_dir ./results/inference
/home/aistudio/work
/opt/conda/envs/python35-paddle120-env/lib/python3.7/site-packages/setuptools/depends.py:2: DeprecationWarning: the imp module is deprecated in favour of importlib; see the module's documentation for alternative uses
  import imp
Model A.D. Param#: Tensor(shape=[1], dtype=int64, place=Place(cpu), stop_gradient=False,
       [4896728])
inference model has been saved into ./results/inference
# 基于推理引擎的模型预测
%cd /home/aistudio/work
!python deploy/infer.py  -cfg configs/resnet18.yaml --save_inference_dir ./results/inference --use_gpu True --image_path images/bottle_good.png
/home/aistudio/work
/opt/conda/envs/python35-paddle120-env/lib/python3.7/site-packages/setuptools/depends.py:2: DeprecationWarning: the imp module is deprecated in favour of importlib; see the module's documentation for alternative uses
  import imp
E0922 16:12:02.416582  4068 analysis_config.cc:95] Please compile with gpu to EnableGpu()
Normal - score:  0.275

* 预测核心代码展示

 # preprocess
    img = inference_engine.preprocess(args.image_path)

    if args.benchmark:
        autolog.times.stamp()

    output = inference_engine.run(img)

    if args.benchmark:
        autolog.times.stamp()

    # postprocess
    output = inference_engine.postprocess(output)
    
    image_score = np.max(output)
    if image_score > args.image_threshold:
        print('Anomaly - score:  {:.3f}'.format(image_score))
    else:
        print('Normal - score:  {:.3f}'.format(image_score))
    output = gaussian_filter(output, sigma=6)[0]
    predict_map = (output > args.pixel_threshold).astype(np.float32)
    save_image = np.concatenate((output, predict_map), axis = 0) * 255

    cv2.imwrite('./output/lele.jpg', save_image.astype(np.uint8))

6、自动化测试脚本

自动化测试为论文复现赛所需,可以自动验证项目的训练测试以及推理的正确性。

%cd /home/aistudio/work
!bash test_tipc/test_train_inference_python.sh test_tipc/configs/fastflow/train_infer_python.txt lite_train_lite_infer 
/home/aistudio/work
W0922 16:37:38.155710  2445 gpu_resources.cc:61] Please NOTE: device: 0, GPU Compute Capability: 7.0, Driver API Version: 11.2, Runtime API Version: 11.2
W0922 16:37:38.159940  2445 gpu_resources.cc:91] device: 0, cuDNN Version: 8.2.
Model A.D. Param#: Tensor(shape=[1], dtype=int64, place=Place(gpu:0), stop_gradient=False,
       [4896728])
Loading training data
Took 0.0018277168273925781
Start training
100%|█| 1/1 [00:02<00:00,  2.34s/it, bottle  Epoch 1 - iter:1/1 - lr:0.001 - los
EVAL:  image-auc : 0.719000 pixel-auc:0.799000
100%|█| 1/1 [00:00<00:00,  1.39it/s, bottle  Epoch 2 - iter:1/1 - lr:0.001 - los
EVAL:  image-auc : 0.578000 pixel-auc:0.800000
100%|█| 1/1 [00:00<00:00,  1.41it/s, bottle  Epoch 3 - iter:1/1 - lr:0.001 - los
EVAL:  image-auc : 0.578000 pixel-auc:0.815000
Training time 0:00:10
save:lite_train_lite_infer
bottle           pixel auroc:0.799    image auroc:0.719
mean             pixel auroc:0.7990    image auroc:0.7190
[33m Run successfully with command - python3.7 train.py --category=bottle --config='configs/resnet18.yaml' --exp_dir=lite_train_lite_infer --epochs=3 --batch_size=32 --data=lite_data!  [0m
W0922 16:37:53.190227  2570 gpu_resources.cc:61] Please NOTE: device: 0, GPU Compute Capability: 7.0, Driver API Version: 11.2, Runtime API Version: 11.2
W0922 16:37:53.194566  2570 gpu_resources.cc:91] device: 0, cuDNN Version: 8.2.
Model A.D. Param#: Tensor(shape=[1], dtype=int64, place=Place(gpu:0), stop_gradient=False,
       [4896728])
EVAL:  image-auc : 0.719000 pixel-auc:0.799000
bottle           pixel auroc:0.799   image auroc:0.719
all              pixel auroc:0.799   image auroc:0.719
[33m Run successfully with command - python3.7 eval.py --category=bottle  --data=lite_data  --exp_dir=lite_train_lite_infer !  [0m
Model A.D. Param#: Tensor(shape=[1], dtype=int64, place=Place(cpu), stop_gradient=False,
       [4896728])
inference model has been saved into ./test_tipc/output/fastflow/lite_train_lite_infer/norm_train_gpus_0
[33m Run successfully with command - python3.7 deploy/export_model.py --category=bottle --exp_dir=lite_train_lite_infer --save_inference_dir=./test_tipc/output/fastflow/lite_train_lite_infer/norm_train_gpus_0!  [0m
Normal - score:  0.298
[33m Run successfully with command - python3.7 deploy/infer.py --use_gpu=True --save_inference_dir=./test_tipc/output/fastflow/lite_train_lite_infer/norm_train_gpus_0 --batch_size=1   --benchmark=False --image_path=images/bottle_good.png > ./test_tipc/output/fastflow/lite_train_lite_infer/python_infer_gpu_batchsize_1.log 2>&1 !  [0m
Normal - score:  0.298
[33m Run successfully with command - python3.7 deploy/infer.py --use_gpu=False --save_inference_dir=./test_tipc/output/fastflow/lite_train_lite_infer/norm_train_gpus_0 --batch_size=1   --benchmark=False --image_path=images/bottle_good.png > ./test_tipc/output/fastflow/lite_train_lite_infer/python_infer_cpu_batchsize_1.log 2>&1 !  [0m

五、代码结构

    |--images                            # 测试使用的样例图片,两张
    |--data                                 # 训练和测试数据集
    |--lite_data                          # 自建立的小数据集,含有bottle 
    |--deploy                            # 预测部署相关
        |--export_model.py        # 导出模型
        |--infer.py                       # 部署预测
    |--results                            # 保存权重和日志 
    |--configs                           # 配置
    |--models                           # 模型实现文件
    |--datasets                         # 数据集加载
    |--utils                                # 工具代码
    |--test_tipc                         # tipc代码
    |--predict.py                      # 预测代码
    |--eval.py                           # 评估代码
    |--train.py                          # 训练代码
    |--train.sh                          # 训练所有类别并进行测试
    |--README.md                 # 用户手册

六、复现总结

FastFlow是一种基于normalizing flow的可移植性很强的,适用于各种异常检测、缺陷检测场景。感谢作者为我们贡献了这一方便有效的方法!

此文章为搬运
原项目链接

Logo

学大模型,用大模型上飞桨星河社区!每天8点V100G算力免费领!免费领取ERNIE 4.0 100w Token >>>

更多推荐