PyTorch使用ResNet18提取图像特征并进行相似度计算
模型部分我参考的是https://blog.csdn.net/sunqiande88/article/details/80100891这篇文章，同样是在Cifar-10上训练。
一、不使用PyTorch中的预训练模型
将训练的模型保存下来接后面使用，保存方式：
torch.save(net.state_dict(), 'path')
加载方式
model = ResNet18()
model.load_state_dict(torch.load('path'))
model.eval() # 测试时候要加上
由于不是使用预训练模型所以输出特征层还是很好输出的，只需要将输出从fc层改为前一层即可：
def forward(self, x):
out = self.conv1(x)
out = self.layer1(out)
out = self.layer2(out)
out = self.layer3(out)
out = self.layer4(out)
out = F.avg_pool2d(out, 4)
out = out.view(out.size(0), -1)
features = out.detach()
out = self.fc(out)
return features # [1,512]
保存ONNX模型：
model = ResNet18()
model.load_state_dict(torch.load(from_path))
model.eval()
dummy_input = torch.randn(1, 3, 32, 32, requires_grad=True)
torch.onnx.export(model, dummy_input, to_path, export_params=True, opset_version=10, do_constant_folding=True
, input_names=["input"], output_names=['output'])
二、使用PyTorch中的预训练模型
模型的保存与加载与上面相同，但是由于使用预训练模型无法修改输出，所以需要使用其他方式修改模型输出。
将训练的模型保存下来接后面使用，保存方式：
torch.save(net.state_dict(), 'path')
加载方式：
model = models.resnet18(pretrained=False)
model.fc = nn.Linear(512, 10)
model.load_state_dict(torch.load(from_path))
model.eval()
# 获取全连接层的输入
class Identity(nn.Module):
def __init__(self):
super(Identity, self).__init__()
def forward(self, x):
return x
model = models.resnet18(pretrained=False)
model.fc = Identity()
x = torch.randn(1, 3, 32, 32)
output = model(x)
保存ONNX模型：
# 初始化并加载模型
model = models.resnet18(pretrained=False)
model.fc = nn.Linear(512, 10)
model.load_state_dict(torch.load('path'))
# 修改fc层输出
model.fc = Identity() # 修改FC层输出
# 调试模式
model.eval()
# 输出到onnx
dummy_input = torch.randn(1, 3, 32, 32, requires_grad=True)
torch.onnx.export(model, dummy_input, to_path, export_params=True, opset_version=10, do_constant_folding=True
, input_names=['input'], output_names=['output'])
三、测试
测试.pth模型部分代码
# 加载图像
def load_image(img_path, transform=None):
imgs = []
for name in sorted(os.listdir(img_path)):
img = Image.open(img_path + name).convert('RGB')
if transform is not None:
img = transform(img)
else:
img = transforms.ToTensor()(img)
imgs.append(img)
return imgs
def predict(imgs):
model = ResNet18()
model.load_state_dict(torch.load('path'))
model.to(device)
model.eval()
imgs = torch.stack(imgs, 0).to(device)
with torch.no_grad():
predicts = model(imgs)
print(predicted)
return predicted
测试ONNX部分分为以下几种方式，
1、使用Python中的onnxruntime库
session = onnxruntime.InferenceSession('ONNX Path')
inputs = {session.get_inputs()[0].name: img.numpy()}
outs = session.run(None, inputs)
return outs
2、使用Python中OpenCV（Version：4.2.0）中的接口
# 加载图像
def load_image_cv(img_path, mean=None, std=None):
img = cv2.imread(img_path, cv2.IMREAD_ANYCOLOR)
if img.shape[2] > 1:
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_ = (img / 255. - mean) / std # 这里根据训练时的格式进行处理
return img_.astype(np.float32)
# 使用opencv运行onnx
def run_onnx_cv(from_path, img):
net = cv2.dnn.readNetFromONNX(from_path)
input = cv2.dnn.blobFromImage(img)
net.setInput(input)
outs = net.forward()
return outs
3、使用C++版本的OpenCV（Version：4.2.0）
//加载图像
cv::Mat load_image_cv(const std::string& fileName, cv::Scalar mean, cv::Scalar std)
{
cv::Mat img = cv::imread(fileName, cv::IMREAD_ANYCOLOR);
if (img.empty()) {
return cv::Mat();
}
img.convertTo(img, CV_32F, 1 / 255.);
cv::subtract(img, mean, img);
cv::divide(img, std, img);
return img;
}
cv::Mat img = load_image_cv("path", cv::Scalar(0.4914, 0.4822, 0.4465), cv::Scalar(0.2023, 0.1994, 0.2010));
cv::dnn::Net net = cv::dnn::readNetFromONNX("ONNX Path");
net.setPreferableBackend(cv::dnn::Backend::DNN_BACKEND_OPENCV);
net.setPreferableTarget(cv::dnn::Target::DNN_TARGET_CPU);
cv::Mat input = cv::dnn::blobFromImage(img);
net.setInput(input);
cv::Mat predicted = net.forward();
效果图片
