Browse Source

OpenVINO metadata fix (#7952)

* Rename OpenVINO meta.yaml to model name

* Rename OpenVINO meta.yaml to model name

* Rename OpenVINO meta.yaml to model name

* fix
modifyDataloader
Glenn Jocher GitHub 2 years ago
parent
commit
0dd66e2dc7
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 15 additions and 14 deletions
  1. +1
    -1
      export.py
  2. +14
    -13
      models/common.py

+ 1
- 1
export.py View File



cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} --data_type {'FP16' if half else 'FP32'}" cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} --data_type {'FP16' if half else 'FP32'}"
subprocess.check_output(cmd.split()) # export subprocess.check_output(cmd.split()) # export
with open(Path(f) / 'meta.yaml', 'w') as g:
with open(Path(f) / file.with_suffix('.yaml'), 'w') as g:
yaml.dump({'stride': int(max(model.stride)), 'names': model.names}, g) # add metadata.yaml yaml.dump({'stride': int(max(model.stride)), 'names': model.names}, g) # add metadata.yaml


LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')

+ 14
- 13
models/common.py View File

super().__init__() super().__init__()
w = str(weights[0] if isinstance(weights, list) else weights) w = str(weights[0] if isinstance(weights, list) else weights)
pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs = self.model_type(w) # get backend pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs = self.model_type(w) # get backend
stride, names = 32, [f'class{i}' for i in range(1000)] # assign defaults
w = attempt_download(w) # download if not local w = attempt_download(w) # download if not local
fp16 &= (pt or jit or onnx or engine) and device.type != 'cpu' # FP16 fp16 &= (pt or jit or onnx or engine) and device.type != 'cpu' # FP16
stride, names = 32, [f'class{i}' for i in range(1000)] # assign defaults
if data: # assign class names (optional)
with open(data, errors='ignore') as f:
names = yaml.safe_load(f)['names']


if pt: # PyTorch if pt: # PyTorch
model = attempt_load(weights if isinstance(weights, list) else w, map_location=device) model = attempt_load(weights if isinstance(weights, list) else w, map_location=device)
network = ie.read_model(model=w, weights=Path(w).with_suffix('.bin')) network = ie.read_model(model=w, weights=Path(w).with_suffix('.bin'))
executable_network = ie.compile_model(model=network, device_name="CPU") executable_network = ie.compile_model(model=network, device_name="CPU")
output_layer = next(iter(executable_network.outputs)) output_layer = next(iter(executable_network.outputs))
self._load_metadata(w.parent / 'meta.yaml') # load metadata
meta = w.with_suffix('.yaml')
if meta.exists():
stride, names = self._load_metadata(meta) # load metadata
elif engine: # TensorRT elif engine: # TensorRT
LOGGER.info(f'Loading {w} for TensorRT inference...') LOGGER.info(f'Loading {w} for TensorRT inference...')
import tensorrt as trt # https://developer.nvidia.com/nvidia-tensorrt-download import tensorrt as trt # https://developer.nvidia.com/nvidia-tensorrt-download
output_details = interpreter.get_output_details() # outputs output_details = interpreter.get_output_details() # outputs
elif tfjs: elif tfjs:
raise Exception('ERROR: YOLOv5 TF.js inference is not supported') raise Exception('ERROR: YOLOv5 TF.js inference is not supported')

self.__dict__.update(locals()) # assign all variables to self self.__dict__.update(locals()) # assign all variables to self
if not hasattr(self, 'names') and data: # assign class names (optional)
with open(data, errors='ignore') as f:
names = yaml.safe_load(f)['names']


def forward(self, im, augment=False, visualize=False, val=False): def forward(self, im, augment=False, visualize=False, val=False):
# YOLOv5 MultiBackend inference # YOLOv5 MultiBackend inference
y = torch.tensor(y, device=self.device) y = torch.tensor(y, device=self.device)
return (y, []) if val else y return (y, []) if val else y


def _load_metadata(self, f='path/to/meta.yaml'):
# Load metadata from meta.yaml if it exists
if Path(f).is_file():
with open(f, errors='ignore') as f:
for k, v in yaml.safe_load(f).items():
setattr(self, k, v) # assign stride, names

def warmup(self, imgsz=(1, 3, 640, 640)): def warmup(self, imgsz=(1, 3, 640, 640)):
# Warmup model by running inference once # Warmup model by running inference once
warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb
tflite &= not edgetpu # *.tflite tflite &= not edgetpu # *.tflite
return pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs return pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs


@staticmethod
def _load_metadata(f='path/to/meta.yaml'):
# Load metadata from meta.yaml if it exists
with open(f, errors='ignore') as f:
d = yaml.safe_load(f)
return d['stride'], d['names'] # assign stride, names



class AutoShape(nn.Module): class AutoShape(nn.Module):
# YOLOv5 input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS # YOLOv5 input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS

Loading…
Cancel
Save