图像特征提取中的迁移学习应用

在大模型训练中，图像特征提取是关键环节。迁移学习为解决数据不足问题提供了有效方案。本文将分享基于预训练模型的图像特征提取实践。

核心思路

利用ImageNet预训练的ResNet50模型，冻结前面层，仅训练最后几层分类器，实现快速特征提取。

实现步骤

环境准备：使用PyTorch框架

import torch
import torchvision.models as models
from torchvision import transforms
from PIL import Image

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

模型加载与修改：

# 加载预训练模型
model = models.resnet50(pretrained=True)

# 冻结前面层
for param in model.parameters():
    param.requires_grad = False

# 修改最后的分类层为自定义输出数
num_classes = 10
model.fc = torch.nn.Linear(model.fc.in_features, num_classes)

特征提取函数：

model.eval()
model.to(device)

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

def extract_features(image_path):
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0).to(device)
    with torch.no_grad():
        features = model(image)
    return features.cpu().numpy()

通过此方法可快速提取图像特征，为后续大模型训练提供高质量输入数据。