数据生成脚本适配样例
==========

import sys
import os
import json
import logging
from functools import reduce
from dataclasses import dataclass
import numpy as np
from numpy import array
DTYPE = {
    "float16": np.float16
}
atol, rtol = 0.005, 0.005
def byte\_2\_bit(drop\_mask, size):
    print(len(drop\_mask))
    drop\_mask\_bit = np.zeros((int(size / 8),), dtype=np.uint8)
    for i in range(int(size / 8)):
        tmp = 0
        for j in range(8):
            tmp += drop\_mask\[i \* 8 + j\] \* (2 \*\* j)
        drop\_mask\_bit\[i\] = tmp
    return drop\_mask\_bit
def get\_input\_attrs\_content(input\_json: str) -> dict:
    content = {}
    input\_json = os.path.realpath(input\_json)
    with open(input\_json, "r") as in\_hadle:
        try:
            content = json.load(in\_hadle)
        except json.decoder.JSONDecodeError:
            logging.error("json.decoder.JSONDecodeError")
        finally:
            pass
    return content
def softmax(x, axis=None):
    x = x.astype(np.float32)
    x\_max = x.max(axis=-1, keepdims=True)
    x\_sub = x - x\_max
    y = np.exp(x\_sub)
    x\_sum = y.sum(axis=-1, keepdims=True)
    ans = y / x\_sum
    ans = ans.astype(np.float16)
    x\_max = x\_max.astype(np.float16)
    x\_sum = x\_sum.astype(np.float16)
    return ans, x\_max, x\_sum
def attention\_score(shape\_info\_tuple,
                    input\_path,
                    output\_path,
                    need\_new):
    qkv\_shape, drop\_mask\_shape, keep\_prob, input\_dtype, output\_dtype = shape\_info\_tuple
    if need\_new:
        q = np.random.uniform(-0.1, 0.1, qkv\_shape).astype(input\_dtype)
        q\_trans = q.transpose(0, 2, 1, 3)
        q\_trans.tofile(input\_path\["query"\])
        k = np.random.uniform(-0.1, 0.1, qkv\_shape).astype(input\_dtype)
        k\_trans = k.transpose(0, 2, 1, 3)
        k\_trans.tofile(input\_path\["key"\])
        v = np.random.uniform(-0.1, 0.1, qkv\_shape).astype(input\_dtype)
        v\_trans = v.transpose(0, 2, 1, 3)
        v\_trans.tofile(input\_path\["value"\])
        if drop\_mask\_shape is not None:
            drop\_mask = np.random.uniform(0, 2, drop\_mask\_shape).astype(np.uint8)
            size = reduce(lambda x, y: x \* y, drop\_mask\_shape)
            drop\_mask\_bit = byte\_2\_bit(
                drop\_mask, int(size)).reshape(qkv\_shape\[0\], qkv\_shape\[1\], qkv\_shape\[2\], int(qkv\_shape\[2\] / 8))
            drop\_mask\_bit.tofile(input\_path\["drop\_mask"\])
    else:
        q = np.fromfile(input\_path\["query"\], input\_dtype).reshape(qkv\_shape).astype(input\_dtype)
        k = np.fromfile(input\_path\["key"\], input\_dtype).reshape(qkv\_shape).astype(input\_dtype)
        v = np.fromfile(input\_path\["value"\], input\_dtype).reshape(qkv\_shape).astype(input\_dtype)
        drop\_mask = np.fromfile(
            input\_path\["drop\_mask"\],
            dtype=np.uint8).reshape(qkv\_shape\[0\], qkv\_shape\[1\], qkv\_shape\[2\], int(qkv\_shape\[2\] / 8))
    qk = np.matmul(q, k.transpose(0, 1, 3, 2))
    softmax\_res, x\_max, x\_sum = softmax(qk)
    drop\_res = 1 / keep\_prob \* np.multiply(
        softmax\_res, drop\_mask.reshape(qkv\_shape\[0\], qkv\_shape\[1\], qkv\_shape\[2\], qkv\_shape\[2\]))
    y = np.matmul(drop\_res, v)
    y = y.transpose(0, 2, 1, 3)
    y.astype(output\_dtype).tofile(output\_path\["attention\_out"\])
    return
def gen\_total\_golden\_output(shape\_info, input\_path, golden\_path, need\_new):
    b\_value, head\_num, s\_value, h\_value = shape\_info.shape\_info
    qkv\_shape = \[b\_value, head\_num, s\_value, h\_value // head\_num\]
    drop\_shape = \[b\_value \* head\_num \* s\_value \* s\_value\]
    attention\_score((qkv\_shape,
                     drop\_shape,
                     shape\_info.keep\_prob,
                     shape\_info.input\_dtype,
                     shape\_info.output\_dtype),
                    input\_path,
                    golden\_path,
                    need\_new)
def get\_data\_file(res\_path, contents):
    data\_files = {}
    for content in contents:
        if content\["data\_file"\] is None:
            continue
        data\_files\[content\["name"\]\] = os.path.join(res\_path, content\["data\_file"\])
    return data\_files
@dataclass
class InputOutputShapeInfo:
    shape\_info: tuple
    keep\_prob: int
    input\_dtype: np.dtype
    output\_dtype: np.dtype
def get\_args():
    """
    需要开发者按需配置
    sys.argv\[1\]: input json file path
    sys.argv\[2\]: output file path
    sys.argv\[3\]: gen data type: input/output/all
    """
    logging.info(f"数据保存路径为：{sys.argv\[2\]}")
    input\_file = sys.argv\[1\]
    res\_path = sys.argv\[2\]
    content = get\_input\_attrs\_content(input\_file)
    input\_files = get\_data\_file(res\_path, content.get("inputs"))
    output\_files = get\_data\_file(res\_path, content.get("outputs"))
    attrs = content.get("attrs")
    cmp\_idx = 0
    for attr in attrs:
        if attr\["name"\] == "input\_layout":
            input\_layout = attr\["value"\]
            logging.info(f"get input\_layout:{input\_layout}")
        if attr\["name"\] == "keep\_prob":
            keep\_prob = attr\["value"\]
            logging.info(f"get keep\_prob:{keep\_prob}")
        if attr\["name"\] == "head\_num":
            head\_num = attr\["value"\]
            logging.info(f"get head\_num:{head\_num}")
    shape = content.get("inputs")\[cmp\_idx\].get("shape")
    b\_value = shape\[input\_layout.find("B")\]
    s\_value = shape\[input\_layout.find("S")\]
    h\_value = shape\[input\_layout.find("H")\]
    input\_dtype = DTYPE.get(content.get("inputs")\[cmp\_idx\].get("dtype"), np.float16)
    output\_dtype = DTYPE.get(content.get("outputs")\[cmp\_idx\].get("dtype"), np.float16)
    shape\_info = InputOutputShapeInfo(shape\_info=(b\_value, head\_num, s\_value, h\_value),
                                      keep\_prob=keep\_prob,
                                      input\_dtype=input\_dtype,
                                      output\_dtype=output\_dtype)
    return shape\_info, input\_files, output\_files
if \_\_name\_\_ == '\_\_main\_\_':
    """
    sys.argv\[1\]: input data file path
    sys.argv\[2\]: output file path
    sys.argv\[3\]: gen data type: input/output/all
    """
    tensor\_shape\_info, input\_file\_list, output\_file\_list = get\_args()
    gen\_total\_golden\_output(tensor\_shape\_info, input\_file\_list, output\_file\_list, need\_new=True)