我有属于你的密码 feature_extractor.py
以下是此文件夹的一部分:
import torch
import torchvision.transforms as transforms
import numpy as np
import cv2
from .model import Net
class Extractor(object):
def __init__(self, model_path, use_cuda=True):
self.net = Net(reid=True)
self.device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu"
state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)['net_dict']
self.net.load_state_dict(state_dict)
print("Loading weights from {}... Done!".format(model_path))
self.net.to(self.device)
self.size = (64, 128)
self.norm = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
def _preprocess(self, im_crops):
def _resize(im, size):
return cv2.resize(im.astype(np.float32) / 255., size)
im_batch = torch.cat([self.norm(_resize(im, self.size)).unsqueeze(0) for im in im_crops], dim=0).float()
return im_batch
def __call__(self, im_crops):
im_batch = self._preprocess(im_crops)
with torch.no_grad():
im_batch = im_batch.to(self.device)
features = self.net(im_batch)
return features.cpu().numpy()
if __name__ == '__main__':
img = cv2.imread("demo.jpg")[:, :, (2, 1, 0)]
extr = Extractor("checkpoint/ckpt.t7")
feature = extr(img)
print(feature.shape)
现在假设有200个请求排成一行继续。为每个请求加载模型的过程使代码运行缓慢。
所以我觉得把pytorch模型放在缓存里可能是个好主意。我把它改成这样:
from redis import Redis
import msgpack as msg
r = Redis('111.222.333.444')
class Extractor(object):
def __init__(self, model_path, use_cuda=True):
try:
self.net = msg.unpackb(r.get('REID_CKPT'))
finally:
self.net = Net(reid=True)
self.device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu"
state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)['net_dict']
self.net.load_state_dict(state_dict)
print("Loading weights from {}... Done!".format(model_path))
self.net.to(self.device)
packed_net = msg.packb(self.net)
r.set('REID_CKPT', packed_net)
self.size = (64, 128)
self.norm = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
不幸的是,出现了以下错误:
File "msgpack/_packer.pyx", line 286, in msgpack._cmsgpack.Packer.pack
File "msgpack/_packer.pyx", line 292, in msgpack._cmsgpack.Packer.pack
File "msgpack/_packer.pyx", line 289, in msgpack._cmsgpack.Packer.pack
File "msgpack/_packer.pyx", line 283, in msgpack._cmsgpack.Packer._pack
TypeError: can not serialize 'Net' object
原因很明显是因为它不能转换net对象( pytorch nn.Module
类)到字节。
如何有效地将pytorch模型保存在缓存中(或者以某种方式将其保存在ram中)并为每个请求调用它?
谢谢大家。
1条答案
按热度按时间2w3kk1z51#
如果您只需要在ram上保持模型状态,那么redis是不必要的。您可以将ram装载为虚拟磁盘,并将模型状态存储在那里。退房
tmpfs
.