表面缺陷检测是一项重要的视觉质量检测任务，是在制造行业和品控中广泛应用的检测方法。它可以通过图像处理和计算机视觉方法来检测生产过程中出现的表面缺陷，例如异物、瑕疵和凸起等。在本文中，我们将介绍如何使用 PyTorch 框架构建表面缺陷检测模型。

PyTorch 是一个基于 Python 的科学计算库，它是一个高度灵活的框架，适用于各种深度学习任务，包括计算机视觉。我们可以使用 PyTorch 中的内置函数和工具来方便地构建表面缺陷检测模型。

首先，我们需要加载图像数据，这里我们使用 PyTorch 中内置的数据集 torchvision.datasets 中的 MNIST 数据集作为示例。 MNIST 数据集包含了大量手写数字的图像，是入门级的图像分类任务。我们将其用于表面缺陷检测任务。

from torchvision import datasets
from torchvision.transforms import ToTensor


# 加载 MNIST 数据集
mnist_train = datasets.MNIST(root="data", train=True, download=True, transform=ToTensor())
mnist_test = datasets.MNIST(root="data", train=False, download=True, transform=ToTensor())

# 显示 MNIST 图像
import matplotlib.pyplot as plt

image, label = mnist_train[0]
plt.imshow(image[0],cmap='gray')
plt.show()

接着，我们将使用 PyTorch 中的卷积神经网络（CNN）来构建表面缺陷检测模型。CNN 在图像处理任务中非常流行，因为它能够自动学习图像特征。在本例中，我们将在 CNN 中使用两个卷积层和两个全连接层，以便对 MNIST 数据集中的手写数字进行分类。这些层将分别使用 ReLU 和 Softmax 激活函数。

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 卷积层
        self.conv1 = nn.Conv2d(1, 6, 5)