wehub-resource-sync
144 lines
5.6 KiB
Python
144 lines
5.6 KiB
Python
import numpy
|
|
import json
|
|
def write_quant_header(file, ic, oc, quant_bit):
|
|
dim_num = file.write(b'\x02')
|
|
shape_dtype = numpy.int16
|
|
if oc > 65535 or ic > 65535:
|
|
shape_dtype = numpy.int32
|
|
dim_length = file.write(numpy.array([oc, ic]).astype(shape_dtype))
|
|
offset = 1 << (quant_bit - 1)
|
|
weight_map = [i for i in range(-offset, offset)]
|
|
if len(weight_map) == 256:
|
|
weight_map.insert(0, 0)
|
|
else:
|
|
weight_map.insert(0, len(weight_map))
|
|
map_length = file.write(numpy.array(weight_map, dtype=numpy.int8))
|
|
header_length = dim_num + dim_length + map_length
|
|
return header_length, shape_dtype == numpy.int32
|
|
|
|
|
|
def repack_low_bits(x, iNeedBits, block_size):
|
|
v = []
|
|
block_number = x.shape[0]
|
|
count = block_size * iNeedBits // 8
|
|
for i in range(0, count):
|
|
v.append(numpy.zeros([block_number, 1]).astype(numpy.uint8))
|
|
iOffset = 0
|
|
cMask = (1 << iNeedBits) - 1
|
|
index = 0
|
|
for i in range(0, block_size):
|
|
p0 = x[:, i:i+1]
|
|
uShift = 8 - iNeedBits - (iOffset % 8)
|
|
if uShift < 0:
|
|
v[index+iOffset // 8] |= ((p0 & cMask) >> (0 - uShift))
|
|
v[index+(iOffset // 8) + 1] |= ((p0 & cMask) << (8 + uShift))
|
|
else:
|
|
v[index+iOffset // 8] |= ((p0 & cMask) << uShift)
|
|
iOffset += iNeedBits
|
|
if iOffset % 8 == 0:
|
|
index += iOffset // 8
|
|
iOffset = 0
|
|
return numpy.concatenate(v, axis=1)
|
|
|
|
class Block:
|
|
def __init__(self):
|
|
self.conv = []
|
|
self.layernorm = []
|
|
|
|
def load_mnn(filename):
|
|
mnn = {}
|
|
with open(filename) as f:
|
|
mnn = json.load(f)
|
|
conv_indexes = []
|
|
layernorm_indexes = []
|
|
blockops = []
|
|
for op in mnn["oplists"]:
|
|
if op['type'] == 'LayerNorm':
|
|
if 'external' in op['main']:
|
|
del op['main']['external']
|
|
if 'gamma' in op['main']:
|
|
del op['main']['gamma']
|
|
if 'beta' in op['main']:
|
|
del op['main']['beta']
|
|
layernorm_indexes.append(len(blockops))
|
|
blockops.append(op)
|
|
continue
|
|
if op['type'] == 'Convolution':
|
|
conv_indexes.append(len(blockops))
|
|
blockops.append(op)
|
|
block = None
|
|
blockes = []
|
|
conv_order = ['attn_q', 'attn_k', 'attn_v', 'attn_output', 'ffn_gate', 'ffn_up', 'ffn_down']
|
|
blockNumber = len(conv_indexes) // len(conv_order)
|
|
print("Layers number: ", blockNumber, ", conv number: ", len(conv_indexes), ", layernorm number:", len(layernorm_indexes))
|
|
block_layernorms = len(layernorm_indexes) // blockNumber
|
|
assert(len(layernorm_indexes) == block_layernorms * blockNumber + 1)
|
|
for i in range(0, blockNumber):
|
|
block = Block()
|
|
sta_conv = len(conv_order) * i
|
|
for j in range(0, len(conv_order)):
|
|
index = conv_indexes[sta_conv + j]
|
|
block.conv.append(blockops[index])
|
|
sta_layernorm = block_layernorms * i
|
|
for j in range(0, block_layernorms):
|
|
index = layernorm_indexes[sta_layernorm + j]
|
|
block.layernorm.append(blockops[index])
|
|
blockes.append(block)
|
|
# Last layernorm and lm
|
|
output_norm = blockops[layernorm_indexes[len(layernorm_indexes)-1]]
|
|
lm = blockops[conv_indexes[len(conv_indexes)-1]]
|
|
lm['name'] = 'output'
|
|
opmap = {}
|
|
opmap['output_norm'] = output_norm
|
|
convs = []
|
|
for i in range(0, len(blockes)):
|
|
_block = blockes[i]
|
|
if len(_block.layernorm) == 2:
|
|
opmap['blk.%d' %i + '.attn_norm']= _block.layernorm[0]
|
|
opmap['blk.%d' %i + '.ffn_norm']= _block.layernorm[1]
|
|
elif len(_block.layernorm) == 6:
|
|
names = ['attn_norm', 'attn_q_norm', 'attn_k_norm', 'post_attention_norm', 'ffn_norm', 'post_ffw_norm']
|
|
for j in range(0, len(_block.layernorm)):
|
|
opmap['blk.%d' %i + '.%s' %names[j]]= _block.layernorm[j]
|
|
else:
|
|
assert(False)
|
|
for j in range(0, 7):
|
|
newname = 'blk.%d' %i + '.' + conv_order[j]
|
|
_block.conv[j]['name'] = newname
|
|
convs.append(_block.conv[j])
|
|
convs.append(lm)
|
|
|
|
return mnn, opmap, convs, blockes, block
|
|
|
|
|
|
def write_quant_parameters(quant_bit, asymc, mnn_weight_file, ic, oc, weight_main, scalebias, mnn_weight_offset, need_scale_treat = True, scale_bit = 32):
|
|
conv = {}
|
|
aMin = 0
|
|
readType = 0
|
|
if asymc:
|
|
# Avoid aMin post treat for bias
|
|
offset = -(1 << (quant_bit - 1))
|
|
aMin = 1
|
|
if need_scale_treat:
|
|
scalebias = scalebias.reshape([-1, 2])
|
|
bias = scalebias[:, 0:1]
|
|
scale = scalebias[:, 1:2]
|
|
bias = bias - offset * scale
|
|
scalebias = numpy.concatenate([bias, scale], axis=1).astype(numpy.float32)
|
|
readType = 1
|
|
header_len, shape_int32 = write_quant_header(mnn_weight_file, ic, oc, quant_bit)
|
|
weight_len = mnn_weight_file.write(weight_main.tobytes()) + header_len
|
|
if scale_bit == 16:
|
|
scalebias_bytes = scalebias.astype(numpy.float16).tobytes()
|
|
else:
|
|
scalebias_bytes = scalebias.tobytes()
|
|
alpha_len = mnn_weight_file.write(scalebias_bytes)
|
|
conv['quanParameter'] = {
|
|
"quantScale": 1.0, "scaleIn": 0.0, "scaleOut": 0.0,
|
|
"useInt32": False, "has_scaleInt": False, "shapeInt32": shape_int32,
|
|
"type": 1, "aMaxOrBits": quant_bit, "aMin": aMin, "readType": readType, "weightSize": 0,
|
|
"scaleStorage": "FP16" if scale_bit == 16 else "FP32",
|
|
}
|
|
conv['external'] = [mnn_weight_offset, weight_len, alpha_len, oc * 4, 0]
|
|
mnn_weight_offset += (weight_len + alpha_len)
|
|
return conv, header_len, mnn_weight_offset |