【论文复现赛】FastFlow_Paddle
【兴智杯论文复现赛】基于Paddle复现基于normalizing flow的无监督缺陷检测FastFlow
一、前言
本项目为百度论文复现赛《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 AUC | pixel-level AUC |
---|---|---|
论文 | 97.9 | 97.2 |
复现 | 98.0 | 97.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的可移植性很强的,适用于各种异常检测、缺陷检测场景。感谢作者为我们贡献了这一方便有效的方法!
此文章为搬运
原项目链接
更多推荐
所有评论(0)