【发布时间】：2022-08-23 06:32:43
【问题描述】：
我正在尝试将 Self-Correction-Human-Parsing 转换为 coreml。 我面临的问题也在https://github.com/pytorch/pytorch/issues/52889 上公开 和https://github.com/apple/coremltools/issues/1085。
evaluate.py 文件（供参考）如下所示（在将模型转换为 coreml 之后）：
import coremltools as ct
import os
import argparse
import numpy as np
import torch
import torchvision
from torch.utils import data
from tqdm import tqdm
from PIL import Image as PILImage
import torchvision.transforms as transforms
import torch.backends.cudnn as cudnn
import networks
from datasets.datasets import LIPDataValSet
from utils.miou import compute_mean_ioU
from utils.transforms import BGR2RGB_transform
from utils.transforms import transform_parsing
import onnxruntime
import onnx
def get_arguments():
\"\"\"Parse all the arguments provided from the CLI.
Returns:
A list of parsed arguments.
\"\"\"
parser = argparse.ArgumentParser(description=\"Self Correction for Human Parsing\")
# Network Structure
parser.add_argument(\"--arch\", type=str, default=\'resnet101\')
# Data Preference
parser.add_argument(\"--data-dir\", type=str, default=\'./data/LIP\')
parser.add_argument(\"--batch-size\", type=int, default=1)
parser.add_argument(\"--input-size\", type=str, default=\'473,473\')
parser.add_argument(\"--num-classes\", type=int, default=20)
parser.add_argument(\"--ignore-label\", type=int, default=255)
parser.add_argument(\"--random-mirror\", action=\"store_true\")
parser.add_argument(\"--random-scale\", action=\"store_true\")
# Evaluation Preference
parser.add_argument(\"--log-dir\", type=str, default=\'./log\')
parser.add_argument(\"--model-restore\", type=str, default=\'./log/checkpoint.pth.tar\')
parser.add_argument(\"--gpu\", type=str, default=\'0\', help=\"choose gpu device.\")
parser.add_argument(\"--save-results\", action=\"store_true\", help=\"whether to save the results.\")
parser.add_argument(\"--flip\", action=\"store_true\", help=\"random flip during the test.\")
parser.add_argument(\"--multi-scales\", type=str, default=\'1\', help=\"multiple scales during the test\")
return parser.parse_args()
def get_palette(num_cls):
\"\"\" Returns the color map for visualizing the segmentation mask.
Args:
num_cls: Number of classes
Returns:
The color map
\"\"\"
n = num_cls
palette = [0] * (n * 3)
for j in range(0, n):
lab = j
palette[j * 3 + 0] = 0
palette[j * 3 + 1] = 0
palette[j * 3 + 2] = 0
i = 0
while lab:
palette[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i))
palette[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i))
palette[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i))
i += 1
lab >>= 3
return palette
def multi_scale_testing(model, batch_input_im, crop_size=[473, 473], flip=True, multi_scales=[1]):
flipped_idx = (15, 14, 17, 16, 19, 18)
if len(batch_input_im.shape) > 4:
batch_input_im = batch_input_im.squeeze()
if len(batch_input_im.shape) == 3:
batch_input_im = batch_input_im.unsqueeze(0)
interp = torch.nn.Upsample(size=crop_size, mode=\'bilinear\', align_corners=True)
ms_outputs = []
for s in multi_scales:
interp_im = torch.nn.Upsample(scale_factor=s, mode=\'bilinear\', align_corners=True)
scaled_im = interp_im(batch_input_im)
print(\"Scaled_im:\",type(scaled_im),scaled_im.shape)
# traced_model = torch.jit.trace(model, scaled_im.to(\"cuda:0\"))
# parsing_output = model(scaled_im)
scripted_model = torch.jit.script(model)
print(\"Trying coreml part\")
core_model = ct.convert(scripted_model,inputs=[ct.TensorType(shape=scaled_im.shape)])
core_model.save(\"human_parsing.mlmodel\")
print(\"human parsing model saved!\")
# torch.onnx.export(model,scaled_im.to(\"cuda:0\"),\"human_parsing.onnx\",opset_version=11)
# print(\"Loading onnx model...\")
# onnx_model = onnx.load(\"human_parsing.onnx\")
# onnx.checker.check_model(onnx_model)
# ort_session = onnxruntime.InferenceSession(\"human_parsing.onnx\")
# def to_numpy(tensor):
# return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
# # x = torch.rand(1,3,473,473)
# # ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
# ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(scaled_im)}
# parsing_output = ort_session.run(None, ort_inputs)
# print(\"Parsing output 1:\", len(parsing_output))
parsing_output = parsing_output[0][-1]
parsing_output = np.expand_dims(parsing_output,axis=0)
parsing_output = torch.from_numpy(parsing_output)
print(\"Parsing output 2:\",parsing_output.shape)
output = parsing_output[0]
if flip:
flipped_output = parsing_output[1]
flipped_output[14:20, :, :] = flipped_output[flipped_idx, :, :]
output += flipped_output.flip(dims=[-1])
output *= 0.5
print(\"output:\", output)
print(\"output:\", output.shape)
print(\"output Type:\", type(output))
output = interp(output.unsqueeze(0))
print(\"output unsqueezed:\", output.shape)
ms_outputs.append(output[0])
ms_fused_parsing_output = torch.stack(ms_outputs)
ms_fused_parsing_output = ms_fused_parsing_output.mean(0)
ms_fused_parsing_output = ms_fused_parsing_output.permute(1, 2, 0) # HWC
parsing = torch.argmax(ms_fused_parsing_output, dim=2)
parsing = parsing.data.cpu().numpy()
ms_fused_parsing_output = ms_fused_parsing_output.data.cpu().numpy()
return parsing, ms_fused_parsing_output
def main():
\"\"\"Create the model and start the evaluation process.\"\"\"
args = get_arguments()
multi_scales = [float(i) for i in args.multi_scales.split(\',\')]
gpus = [int(i) for i in args.gpu.split(\',\')]
assert len(gpus) == 1
if not args.gpu == \'None\':
os.environ[\"CUDA_VISIBLE_DEVICES\"] = args.gpu
cudnn.benchmark = True
cudnn.enabled = True
h, w = map(int, args.input_size.split(\',\'))
input_size = [h, w]
model = networks.init_model(args.arch, num_classes=args.num_classes, pretrained=None)
IMAGE_MEAN = model.mean
IMAGE_STD = model.std
INPUT_SPACE = model.input_space
print(\'image mean: {}\'.format(IMAGE_MEAN))
print(\'image std: {}\'.format(IMAGE_STD))
print(\'input space:{}\'.format(INPUT_SPACE))
if INPUT_SPACE == \'BGR\':
print(\'BGR Transformation\')
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=IMAGE_MEAN,
std=IMAGE_STD),
])
if INPUT_SPACE == \'RGB\':
print(\'RGB Transformation\')
transform = transforms.Compose([
transforms.ToTensor(),
BGR2RGB_transform(),
transforms.Normalize(mean=IMAGE_MEAN,
std=IMAGE_STD),
])
# Data loader
lip_test_dataset = LIPDataValSet(args.data_dir, \'val\', crop_size=input_size, transform=transform, flip=args.flip)
num_samples = len(lip_test_dataset)
print(\'Total testing sample numbers: {}\'.format(num_samples))
testloader = data.DataLoader(lip_test_dataset, batch_size=args.batch_size, shuffle=False, pin_memory=True)
# num_samples =1
# Load model weight
state_dict = torch.load(args.model_restore)[\'state_dict\']
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
model.cuda()
model.eval()
sp_results_dir = os.path.join(args.log_dir, \'sp_results\')
if not os.path.exists(sp_results_dir):
os.makedirs(sp_results_dir)
palette = get_palette(20)
parsing_preds = []
scales = np.zeros((num_samples, 2), dtype=np.float32)
centers = np.zeros((num_samples, 2), dtype=np.int32)
with torch.no_grad():
for idx, batch in enumerate(tqdm(testloader)):
image, meta = batch
if (len(image.shape) > 4):
image = image.squeeze()
im_name = meta[\'name\'][0]
c = meta[\'center\'].numpy()[0]
s = meta[\'scale\'].numpy()[0]
w = meta[\'width\'].numpy()[0]
h = meta[\'height\'].numpy()[0]
scales[idx, :] = s
centers[idx, :] = c
parsing, logits = multi_scale_testing(model, image.cuda(), crop_size=input_size, flip=args.flip,
multi_scales=multi_scales)
print(\"Parsing:\",parsing.shape)
print(\"Logits:\", logits.shape)
# if args.save_results:
if True:
print(\"Inside Save_results\")
parsing_result = transform_parsing(parsing, c, s, w, h, input_size)
parsing_result_path = os.path.join(sp_results_dir, im_name + \'.png\')
# print(\"Parsing Result Path:\", parsing_result_path)
output_im = PILImage.fromarray(np.asarray(parsing_result, dtype=np.uint8))
output_im.putpalette(palette)
output_im.save(parsing_result_path)
parsing_preds.append(parsing)
assert len(parsing_preds) == num_samples
mIoU = compute_mean_ioU(parsing_preds, scales, centers, args.num_classes, args.data_dir, input_size)
print(mIoU)
return
if __name__ == \'__main__\':
main()
我已成功将模型移植到 onnx。 但是我在将其转换为 coreml 时遇到了问题。
Traceback (most recent call last):
File \"evaluate.py\", line 262, in <module>
main()
File \"evaluate.py\", line 240, in main
parsing, logits = multi_scale_testing(model, image.cuda(), crop_size=input_size, flip=args.flip,
File \"evaluate.py\", line 102, in multi_scale_testing
core_model = ct.convert(scripted_model,inputs=[ct.TensorType(shape=scaled_im.shape)])
File \"/anaconda/envs/schp/lib/python3.8/site-packages/coremltools/converters/_converters_entry.py\", line 176, in convert
mlmodel = mil_convert(
File \"/anaconda/envs/schp/lib/python3.8/site-packages/coremltools/converters/mil/converter.py\", line 128, in mil_convert
proto = mil_convert_to_proto(model, convert_from, convert_to,
File \"/anaconda/envs/schp/lib/python3.8/site-packages/coremltools/converters/mil/converter.py\", line 171, in mil_convert_to_proto
prog = frontend_converter(model, **kwargs)
File \"/anaconda/envs/schp/lib/python3.8/site-packages/coremltools/converters/mil/converter.py\", line 85, in __call__
return load(*args, **kwargs)
File \"/anaconda/envs/schp/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/load.py\", line 72, in load
converter = TorchConverter(torchscript, inputs, outputs, cut_at_symbols)
File \"/anaconda/envs/schp/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/converter.py\", line 142, in __init__
raw_graph, params_dict = self._expand_and_optimize_ir(self.torchscript)
File \"/anaconda/envs/schp/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/converter.py\", line 250, in _expand_and_optimize_ir
graph, params = _torch._C._jit_pass_lower_graph(
RuntimeError: Unknown type bool encountered in graph lowering. This type is not supported in ONNX export.
标签： python ios pytorch coremltools torchscript