在前两篇博文中,我使用Win11+docker配置了anomalib,并成功的调用了GPU运行了示例程序。这次我准备使用anomalib训练我自己的数据集。
数据集是我在工作中收集到的火腿肠缺陷数据,与MVTec等数据不同,我的火腿肠数据来源于多台设备和多个品种,因此,它们表面的纹理与颜色差异以及不同设备间的成像质量差异要远大于标准数据集,使得检测变得更难了。
我的合格品图像如下图:
我的缺陷品图像如下图:
为了训练这个数据集,我首先将图像拷贝到docker中。使用以下代码可以将本机的文件复制到docker容器。
docker cp path/src 镜像名称:path/dst
这里需要注意的是,我们在docker 中克隆的anomalib代码是在workspace文件夹中的,copy数据集的时候记得加上这一级文件路径。我的路径是workspace\anomalib\datasets\JY
接下来,按照docker文档的提示,Custom Data — Anomalib 2022 documentation,把代码修改为自己的路径。我的数据集是包含了合格样本和异常样本的,但只包含图像标签,并没有标注异常区域的mask,因此我需要使用classification类别的task。
除了路径外,还有几个示例代码中没有设置的参数需要按照自己的数据情况设置下。第一个是normal_split_ratio ,表示将多少个正常样本加入到测试集。我的正常样本有1000个,异常样本有100个,因此我设置其为0.1。另外我训练时显存超出了我的物理显存,总共用了约10G,导致其占用了约3个G的共享显存。我担心这会影响训练速度,就将batchsize改小了一些,默认是32,我将它改成了16,但显存占用并没有变小。我觉得这个batchsize可能是提取特征时用,但patchcore训练的绝大部分时间是在抽取Core Set,或许这个环节不受Batchsize的影响吧。
因为我的图像数量有些多,单个类别有1000个样本,MVTec AD是200多个,这导致我的训练时间很长,约2个小时吧。
训练结果如下。我的分类结果相比于MVTec要低很多,因为我的数据集太复杂了,相比于MVTec数据集,我的数据集中目标的一致性要低很多,我觉得这是影响准确率的主要原因。
文章最后,附上我的训练代码,当然基本都是和文档中一致啦,只修改了很少一点参数。
# Import the datamodule
from anomalib.data import Folder
# Create the datamodule
datamodule = Folder(
name="JY",
root="datasets/JY",
normal_dir="good",
abnormal_dir="ng",
task="classification",
train_batch_size = 16,
eval_batch_size = 16,
num_workers = 0,
normal_split_ratio = 0.1
)
# Setup the datamodule
datamodule.setup()
i, train_data = next(enumerate(datamodule.train_dataloader()))
print(train_data.keys())
# dict_keys(['image_path', 'label', 'image'])
i, val_data = next(enumerate(datamodule.val_dataloader()))
print(val_data.keys())
# dict_keys(['image_path', 'label', 'image'])
i, test_data = next(enumerate(datamodule.test_dataloader()))
print(test_data.keys())
# dict_keys(['image_path', 'label', 'image'])
# Import the model and engine
from anomalib.models import Patchcore
from anomalib.engine import Engine
# Create the model and engine
model = Patchcore()
engine = Engine(task="classification")
# Train a Patchcore model on the given datamodule
engine.train(datamodule=datamodule, model=model)