加载模型进行推理
打算使用训练好的模型提取MS1M的人脸图像的特征进行聚类，记录一下。
模型加载
- 将Pytorch模型转换成ONNX模型的主要函数为
torch.onnx.export() - 对pytorch模型(nn.Module)对象调用
named_parameters()可以得到对应的参数名(key)和参数值(value) - 对pytorch模型对象调用
load_state_dict()可以将预训练模型的参数加载到定义的模型中，注意strict参数
def load_ckpt(path, model, ignores=[], strict=True, optimizer=None, not_match=True):
if os.path.exists(path):
print("Loading checkpoint '{}'".format(path))
ckpt = torch.load(path, map_location=torch.device("cpu"))
# 去掉预训练模型中的module.base.前缀
from collections import OrderedDict
new_ckpt_state_dict = OrderedDict()
if not_match:
for k, v in ckpt["state_dict"].items():
new_ckpt_state_dict[k[12:]] = v
else:
new_ckpt_state_dict = ckpt["state_dict"]
if len(ignores) > 0:
assert optimizer == None
keys = set(new_ckpt_state_dict.keys())
for ignore in ignores:
if ignore in keys:
print("Ignoring {}".format(ignore))
del new_ckpt_state_dict[ignore]
else:
raise ValueError("can not find {} in load_path".foramt(ignore))
model.load_state_dict(new_ckpt_state_dict, strict=strict)
if not strict:
# 如果load_state_dict()中strict参数为False, 那么只会加载两者都有的keys
pretrained_keys = set(new_ckpt_state_dict.keys())
model_keys = set([k for k, _ in model.named_parameters()])
for k in model_keys - pretrained_keys:
print("Warning: {} not loaded\n".format(k))
if optimizer != None:
assert len(ignores) == 0
optimizer.load_state_dict(ckpt["optimizer"])
print("=> Loaded checkpoint '{}' (step {})".format(path, ckpt["epoch"]))
return ckpt["epoch"], ckpt["best_prec1"]
else:
assert False, "=> no checkpoint found at {}".format(path)
使用PyTorch模型进行推理，需要包含模型代码。也可以使用ONNX模型进行推理，这样代码更简洁。下面的代码将PyTorch模型转换成ONNX模型。
def torch2onnx(model_path, output_path):
ckpt = torch.load(model_path, map_location="cpu")
model = resnet50(feature_dim=256)
new_ckpt_state_dict = OrderedDict()
for k, v in ckpt["state_dict"].items():
new_ckpt_state_dict[k[12:]] = v
model.load_state_dict(new_ckpt_state_dict)
dummy_input = torch.randn(1, 3, 112, 112, device="cpu")
input_names = ["input"]
output_names = ["output"]
torch.onnx.export(model, dummy_input, output_path, verbose=True, input_names=input_names, output_names=output_names)
def checkonnx(onnx_model_path):
# Load the ONNX model
onnx_model = onnx.load(onnx_model_path)
# Check that the IR is well formed
onnx.checker.check_model(onnx_model)
ort_session = ort.InferenceSession(onnx_model_path)
outputs = ort_session.run(None, {'input': np.random.randn(1, 3, 112, 112).astype(np.float32)})
print(outputs)
Pytorch模型的保存和加载
参考资料：《模型推理之服务化：Torch模型服务化？难用！》
上面的内容介绍了一种Pytorch模型的加载方法和ONNX模型的加载方法。实际上，Pytorch模型有3种保存和加载方式。
- 1 仅仅保存模型中的参数(上面介绍的那种)
// save
torch.save(model.state_dict(), "model.pkl")
// load
model.load_state_dict(torch.load("model.pkl"))
- 2 保存整个模型
// save
torch.save(model, "model.pt")
// load
model = torch.load("model.pt")
- 3 Pytorch为了部署推出的Torchscript格式
这种格式的Pytorch模型可能很多人比较陌生，其目的是把模型静态化并保存。
// save
script_model = torch.jit.trace(model, dummy_input)
torch.jit.save(script_model, "model.pth")
// load
script_model = torch.jit.load("model.pth")
TorchScript is a way to create serializable and optimizable models from PyTorch code. Any TorchScript program can be saved from a Python process and loaded in a process where there is no Python dependency.
We provide tools to incrementally transition a model from a pure Python program to a TorchScript program that can be run independently from Python, such as in a standalone C++ program. This makes it possible to train models in PyTorch using familiar tools in Python and then export the model via TorchScript to a production environment where Python programs may be disadvantageous for performance and multi-threading reasons.
数据准备
使用opencv读取图像
cv2.imread()读取的图像是BGR格式，tensor的排列为[H, W, C], 每个元素的数据类型为uint8transforms.ToTensor()将tensor的排列变成[C, H, W]，同时每个元素的数据类型为torch.float32，范围归一化成[0, 1]
import cv2
import torchvision.transforms as transforms
img_path = "/Users/tp/Documents/Data/MS1M/emore_ltc/sample_images/10/10_707.jpg"
img = cv2.imread(img_path)
print(img.shape)
>> (112, 112, 3)
print(img.dtype)
>> uint8
# Convert BGR to RGB
img = img[:, :, ::-1]
img_part = cv2.resize(img, (2, 2))
print(img_part.shape)
>> (2, 2, 3)
print(img_part)
>> [[186 206 234]
[ 73 94 112]]
[[152 180 211]
[110 133 151]]]
tensor1 = transforms.ToTensor()(img_part)
print(tensor1)
>> tensor([[[0.7294, 0.2863],
[0.5961, 0.4314]],
[[0.8078, 0.3686],
[0.7059, 0.5216]],
[[0.9176, 0.4392],
[0.8275, 0.5922]]])
print(tensor1.shape)
>> torch.Size([3, 2, 2])
使用PIL的Image读取图像
- Image读取的图像为RGB格式
from PIL import Image
img_pil = Image.open(img_path)
print(type(img_pil))
print(img_pil.size)
>> <class 'PIL.JpegImagePlugin.JpegImageFile'>
(112, 112)
tensor_pil = transforms.ToTensor()(img_pil)
print(tensor_pil[:, :2, :2])
>>
tensor([[[0.2549, 0.2588],
[0.2863, 0.2824]],
[[0.2392, 0.2431],
[0.2706, 0.2667]],
[[0.1961, 0.2000],
[0.2275, 0.2235]]])
img_cv = cv2.imread(img_path)
tensor_cv = transforms.ToTensor()(img_cv)
print(tensor_cv[:, :2, :2])
>>
tensor([[[0.1961, 0.2000],
[0.2275, 0.2235]],
[[0.2392, 0.2431],
[0.2706, 0.2667]],
[[0.2549, 0.2588],
[0.2863, 0.2824]]])
构建InferDataset
class Infer_Dataset(data.Dataset):
def __init__(self, img_dir, transform=None):
self.img_dir = img_dir
self.transform = transform
self.img_list = os.listdir(self.img_dir)
def __getitem__(self, index):
img_path = os.path.join(self.img_dir, self.img_list[index])
img = cv2.imread(img_path)
# 将BGR转换成RGB
img = img[:,:,::-1]
img = cv2.resize(img, (112, 112))
if self.transform is not None:
img = self.transform(img)
else:
img = torch.from_numpy(img)
return img_path, img
def __len__(self):
return len(self.img_list)
模型推理
PyTorch模型推理
feature_dim = 256
model_path = "/Users/tp/Documents/Data/MS1M/resnet50_part0_train.pth.tar"
img_dir ="/Users/tp/Documents/Data/MS1M/emore_ltc/sample_images/10"
model = resnet50(feature_dim=feature_dim)
load_ckpt(model_path, model, strict=False, not_match=True)
normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.3125, 0.3125, 0.3125])
infer_transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
inferset = Infer_Dataset(img_dir, transform=infer_transform)
# num_workers是用来指定开多进程的数量，默认值为0，表示不启用多进程
inferloader = torch.utils.data.DataLoader(inferset, batch_size=16, shuffle=False, num_workers=0, drop_last=False)
model.eval()
for idx, (img_path, img) in enumerate(inferloader):
print(type(img))
embedding = model.forward(img)
print(embedding)
ONNX模型推理
onnx_session = ort.InferenceSession(output_path)
for idx, (img_path, img) in enumerate(inferloader):
print(type(img))
img = np.array(img).astype(np.float32)
# 使用ONNX推理, 输入的shape必须和构建ONNX模型时的shape完全一致
img = img[0:1, :, :, :]
print(img.shape)
embedding = onnx_session.run(None, input_feed={"input": img})
print(embedding)
- PyTorch模型的输入类型自然是torch.tensor()，ONNX模型的输入类型是numpy.array()。
- 关于使用ONNX进行模型推理，可以参考onnx/model中的workflow_scripts代码。一个模块化的处理方式是写一个ONNXModel类。
