web-dev-qa-db-ja.com

pytorchモデルをredisキャッシュに保持して、ビデオストリーミングのモデルに高速にアクセスする方法

ここにこのフォルダの一部である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に保持）して、リクエストごとに呼び出すにはどうすればよいですか？

みんな、ありがとう。

pythoncachingredisvideo-streamingpytorch
4
Masoud Masoumi Moghadam

RAMにモデルの状態を保持するだけでよい場合は、Redisは必要ありません。代わりに、RAMを仮想ディスクとしてマウントし、そこにモデルの状態を保存できます。tmpfsをチェックしてください。

1
roman