88 lines
3.9 KiB
Python
88 lines
3.9 KiB
Python
# Copyright (C) 2022 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# flake8: noqa
|
|
|
|
import numpy as np
|
|
from typing import Union
|
|
from openvino.runtime import Type, Shape
|
|
|
|
|
|
def pack_data(array: np.ndarray, type: Type) -> np.ndarray:
|
|
"""Represent array values as u1,u4 or i4 openvino element type and pack them into uint8 numpy array.
|
|
|
|
If the number of elements in array is odd we pad them with zero value to be able to fit the bit
|
|
sequence into the uint8 array.
|
|
|
|
Example: two uint8 values - [7, 8] can be represented as uint4 values and be packed into one int8
|
|
value - [120], because [7, 8] bit representation is [0111, 1000] will be viewed
|
|
as [01111000], which is bit representation of [120].
|
|
|
|
:param array: numpy array with values to pack.
|
|
:type array: numpy array
|
|
:param type: Type to interpret the array values. Type must be u1, u4 or i4.
|
|
:type type: openvino.runtime.Type
|
|
"""
|
|
assert type in [Type.u1, Type.u4, Type.i4], "Packing algorithm for the" "data types stored in 1, 2 or 4 bits"
|
|
|
|
minimum_regular_dtype = np.int8 if type == Type.i4 else np.uint8
|
|
casted_to_regular_type = array.astype(dtype=minimum_regular_dtype, casting="unsafe")
|
|
if not np.array_equal(casted_to_regular_type, array):
|
|
raise RuntimeError(f'The conversion of array "{array}" to dtype' f' "{casted_to_regular_type}" results in rounding')
|
|
|
|
data_size = casted_to_regular_type.size
|
|
num_bits = type.bitwidth
|
|
|
|
assert num_bits < 8 and 8 % num_bits == 0, "Packing algorithm for the" "data types stored in 1, 2 or 4 bits"
|
|
num_values_fitting_into_uint8 = 8 // num_bits
|
|
pad = (-data_size) % num_values_fitting_into_uint8
|
|
|
|
flattened = casted_to_regular_type.flatten()
|
|
padded = np.concatenate((flattened, np.zeros([pad], dtype=minimum_regular_dtype))) # type: ignore
|
|
assert padded.size % num_values_fitting_into_uint8 == 0
|
|
|
|
bit_order_little = (padded[:, None] & (1 << np.arange(num_bits)) > 0).astype(minimum_regular_dtype)
|
|
bit_order_big = np.flip(bit_order_little, axis=1) # type: ignore
|
|
bit_order_big_flattened = bit_order_big.flatten()
|
|
|
|
return np.packbits(bit_order_big_flattened)
|
|
|
|
|
|
def unpack_data(array: np.ndarray, type: Type, shape: Union[list, Shape]) -> np.ndarray:
|
|
"""Extract openvino element type values from array into new uint8/int8 array given shape.
|
|
|
|
Example: uint8 value [120] can be represented as two u4 values and be unpacked into [7, 8]
|
|
because [120] bit representation is [01111000] will be viewed as [0111, 1000],
|
|
which is bit representation of [7, 8].
|
|
|
|
:param array: numpy array to unpack.
|
|
:type array: numpy array
|
|
:param type: Type to extract from array values. Type must be u1, u4 or i4.
|
|
:type type: openvino.runtime.Type
|
|
:param shape: the new shape for the unpacked array.
|
|
:type shape: Union[list, openvino.runtime.Shape]
|
|
"""
|
|
assert type in [Type.u1, Type.u4, Type.i4], "Unpacking algorithm for the" "data types stored in 1, 2 or 4 bits"
|
|
unpacked = np.unpackbits(array.view(np.uint8))
|
|
shape = list(shape)
|
|
if type.bitwidth == 1:
|
|
return np.resize(unpacked, shape)
|
|
else:
|
|
unpacked = unpacked.reshape(-1, type.bitwidth)
|
|
padding_shape = (unpacked.shape[0], 8 - type.bitwidth)
|
|
padding = np.ndarray(padding_shape, np.uint8) # type: np.ndarray
|
|
if type == Type.i4:
|
|
for axis, bits in enumerate(unpacked):
|
|
if bits[0] == 1:
|
|
padding[axis] = np.ones((padding_shape[1],), np.uint8)
|
|
else:
|
|
padding[axis] = np.zeros((padding_shape[1],), np.uint8)
|
|
else:
|
|
padding = np.zeros(padding_shape, np.uint8)
|
|
padded = np.concatenate((padding, unpacked), 1) # type: ignore
|
|
packed = np.packbits(padded, 1)
|
|
if type == Type.i4:
|
|
return np.resize(packed, shape).astype(dtype=np.int8)
|
|
else:
|
|
return np.resize(packed, shape)
|