使用pytorch中的resnet预训练模型进行快速图像分类

加载模型

import torch
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
# or any of these variants
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet34', pretrained=True)
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet101', pretrained=True)
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet152', pretrained=True)
model.eval()

可供选择的模型有很多，并且不断有新的模型加入，可以参考torchvision官方文档。分类常用的如：[AlexNet, ConvNeXt, DenseNet, EfficientNet, EfficientNetV2, GoogLeNet, Inception V3, MNASNet, MobileNet V2, MobileNet V3, RegNet, ResNet, ResNeXt, ShuffleNet V2, SqueezeNet, SwinTransformer, VGG, VisionTransformer, Wide ResNet]。这里以resnet为例。

处理数据

# sample execution (requires torchvision)
from PIL import Image
from torchvision import transforms
import torchvision
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

path_to_data = '/path/to/data/'
BATCH_SIZE = 256

test_data = torchvision.datasets.ImageFolder(path_to_data, preprocess)
data_loader = torch.utils.data.DataLoader(test_data, batch_size=BATCH_SIZE)

该模型的输入需要设置 batchsize（N），3 通道的 RGB 图像（3 x H x W），H 和 W 最小为224，图像归一化到 [0, 1] 上，均值 mean = [0.485, 0.456, 0.406] 方差 std = [0.229, 0.224, 0.225]。本例中使用 torchvision 的 dataloader 组织数据。由于我们测试时没有标签，图像文件保存在 /path/to/data/ 下的一个文件夹内即可：

/path/to/data/
└── class_x
     ├── xxx.jpg
     ├── xxy.jpg
     └── ...

测试

# load filename of test images
image_names = test_data.samples
# move the input and model to GPU for speed if available
model.to('cuda')

count = 0
result = [] # save the filename if it is predicted to be class 703 in imagenet labels.
with torch.no_grad():
    for (x, y) in tqdm(data_loader, desc="Evaluating", leave=False):
        x = x.to('cuda')
        y = y.to('cuda')

        output = model(x)
        # Tensor of shape 1000, with confidence scores over Imagenet's 1000 classes
        # print(output[0])
        for index in range(output.shape[0]):
            if torch.argmax(output[index,...])==703:
                result.append(image_names[count*BATCH_SIZE+index][0])
        count += 1
        # The output has unnormalized scores. To get probabilities, you can run a softmax on it.
        # print(torch.nn.functional.softmax(output[-1], dim=0))

model 的输出是图像的 logits，1000 维向量，通过查找其中的最大值可以得到其类别，imagenet 对应的类别标注参考 imagenet 或者IMAGENET 1000 Class List。要想得到 [0,1] 归一化后的概率表示，需要再进行一次 softmax。

全部代码如下

我们可以将需要查找的类别图像保存到一个目录（targetdir）下，也可以输出所有图像的路径。

import torch
#model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
# or any of these variants
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet34', pretrained=True)
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet101', pretrained=True)
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet152', pretrained=True)
model.eval()

# sample execution (requires torchvision)
from PIL import Image
from torchvision import transforms
import torchvision
from tqdm import tqdm
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

path_to_data = '/home/llj/ld/data/OpenEdit/train_latest/images/input_image/'
BATCH_SIZE = 256
nThreads = 4
target_dir = 'targetdir'


test_data = torchvision.datasets.ImageFolder(path_to_data, preprocess)
image_names = test_data.samples
data_loader = torch.utils.data.DataLoader(test_data, batch_size=BATCH_SIZE,)
model.to('cuda')

probability = []
count = 0
result = []

with torch.no_grad():
    for (x, y) in tqdm(data_loader, desc="Evaluating", leave=False):
        x = x.to('cuda')
        y = y.to('cuda')

        output = model(x)
        for index in range(output.shape[0]):
            if output[index,703]>0.5:
                result.append(image_names[count*BATCH_SIZE+index][0])
        count += 1
        # Tensor of shape 1000, with confidence scores over Imagenet's 1000 classes
        # print(output[-1])
        # The output has unnormalized scores. To get probabilities, you can run a softmax on it.
        # print(torch.nn.functional.softmax(output[-1], dim=0))
        
from shutil import copyfile
import os

os.makedirs(target_dir, exist_ok=True)

for r in result:
    copyfile(r, target_dir+'/'+r.split('/')[-1])
for r in result:
    print(r)

扩展应用

想通过缩略图找原图？汇报时使用压缩过的图像现在想找原图？如何从大量图像文件中快速找到与目标图像相似的那个？使用pytorch中的resnet预训练模型进行特征提取，以及查找相似图像

参考：

ResNet: Deep residual networks pre-trained on ImageNet
PyTorch Image Classification