fix cpu inference

Link: https://code.alibaba-inc.com/pai-vision/EasyCV/codereview/10267433

    * fix cpu inference
pull/207/head
jiangnana.jnn 2022-09-27 12:30:38 +08:00
parent 7f106ba446
commit 2856af7f42
3 changed files with 46 additions and 36 deletions

View File

@ -292,9 +292,8 @@ class PredictorV2(object):
return collate(inputs, samples_per_gpu=self.batch_size)
def _to_device(self, inputs):
target_gpus = [-1] if self.device == 'cpu' else [
torch.cuda.current_device()
]
target_gpus = [-1] if str(
self.device) == 'cpu' else [torch.cuda.current_device()]
_, kwargs = scatter_kwargs(None, inputs, target_gpus=target_gpus)
return kwargs[0]

View File

@ -2,5 +2,5 @@
# GENERATED VERSION FILE
# TIME: Thu Nov 5 14:17:50 2020
__version__ = '0.6.3.5'
short_version = '0.6.3.5'
__version__ = '0.6.3.6'
short_version = '0.6.3.6'

View File

@ -6,35 +6,25 @@ import tempfile
import unittest
import numpy as np
from mmcv import Config
from PIL import Image
from tests.ut_config import (MODEL_CONFIG_SEGFORMER,
PRETRAINED_MODEL_MASK2FORMER_DIR,
PRETRAINED_MODEL_SEGFORMER, TEST_IMAGES_DIR)
from easycv.file import io
from easycv.predictors.segmentation import SegmentationPredictor
class SegmentationPredictorTest(unittest.TestCase):
img_path = os.path.join(TEST_IMAGES_DIR, '000000289059.jpg')
def setUp(self):
print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
def test_single(self):
segmentation_model_path = PRETRAINED_MODEL_SEGFORMER
segmentation_model_config = MODEL_CONFIG_SEGFORMER
img_path = os.path.join(TEST_IMAGES_DIR, '000000289059.jpg')
img = np.asarray(Image.open(img_path))
predict_pipeline = SegmentationPredictor(
model_path=segmentation_model_path,
config_file=segmentation_model_config)
outputs = predict_pipeline(img_path, keep_inputs=True)
self.assertEqual(len(outputs), 1)
results = outputs[0]
def _assert_results(self, results, img_path):
self.assertEqual(results['inputs'], img_path)
img = np.asarray(Image.open(img_path))
self.assertListEqual(
list(img.shape)[:2], list(results['seg_pred'].shape))
self.assertListEqual(results['seg_pred'][1, :10].tolist(),
@ -42,13 +32,43 @@ class SegmentationPredictorTest(unittest.TestCase):
self.assertListEqual(results['seg_pred'][-1, -10:].tolist(),
[133 for i in range(10)])
def test_single_cpu(self):
cfg_file = MODEL_CONFIG_SEGFORMER
model_path = PRETRAINED_MODEL_SEGFORMER
with tempfile.NamedTemporaryFile(suffix='.py') as tmppath:
with io.open(cfg_file, 'r') as f:
cfg_str = f.read()
cfg_str = cfg_str.replace('SyncBN', 'BN')
with io.open(tmppath.name, 'w') as f:
f.write(cfg_str)
cfg = Config.fromfile(tmppath.name)
predict_pipeline = SegmentationPredictor(
model_path=model_path, config_file=cfg, device='cpu')
outputs = predict_pipeline(self.img_path, keep_inputs=True)
self.assertEqual(len(outputs), 1)
results = outputs[0]
self._assert_results(results, self.img_path)
def test_single(self):
segmentation_model_path = PRETRAINED_MODEL_SEGFORMER
segmentation_model_config = MODEL_CONFIG_SEGFORMER
predict_pipeline = SegmentationPredictor(
model_path=segmentation_model_path,
config_file=segmentation_model_config)
outputs = predict_pipeline(self.img_path, keep_inputs=True)
self.assertEqual(len(outputs), 1)
results = outputs[0]
self._assert_results(results, self.img_path)
def test_batch(self):
segmentation_model_path = PRETRAINED_MODEL_SEGFORMER
segmentation_model_config = MODEL_CONFIG_SEGFORMER
img_path = os.path.join(TEST_IMAGES_DIR, '000000289059.jpg')
img = np.asarray(Image.open(img_path))
predict_pipeline = SegmentationPredictor(
model_path=segmentation_model_path,
config_file=segmentation_model_config,
@ -56,24 +76,15 @@ class SegmentationPredictorTest(unittest.TestCase):
total_samples = 3
outputs = predict_pipeline(
[img_path] * total_samples, keep_inputs=True)
[self.img_path] * total_samples, keep_inputs=True)
self.assertEqual(len(outputs), 3)
for i in range(len(outputs)):
self.assertEqual(outputs[i]['inputs'], img_path)
self.assertListEqual(
list(img.shape)[:2], list(outputs[i]['seg_pred'].shape))
self.assertListEqual(outputs[i]['seg_pred'][1, :10].tolist(),
[161 for i in range(10)])
self.assertListEqual(outputs[i]['seg_pred'][-1, -10:].tolist(),
[133 for i in range(10)])
self._assert_results(outputs[i], self.img_path)
def test_dump(self):
segmentation_model_path = PRETRAINED_MODEL_SEGFORMER
segmentation_model_config = MODEL_CONFIG_SEGFORMER
img_path = os.path.join(TEST_IMAGES_DIR, '000000289059.jpg')
temp_dir = tempfile.TemporaryDirectory().name
if not os.path.exists(temp_dir):
os.makedirs(temp_dir)
@ -88,7 +99,7 @@ class SegmentationPredictorTest(unittest.TestCase):
total_samples = 3
outputs = predict_pipeline(
[img_path] * total_samples, keep_inputs=False)
[self.img_path] * total_samples, keep_inputs=False)
self.assertEqual(outputs, [])
with open(tmp_path, 'rb') as f: