schihei/mixed-precision-neural-networks-on-risc-v-cores
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Adding new features

Browse Source
This commit is contained in:
alexmr09
2024-07-23 13:00:49 +03:00
parent 9e044fd7fc
commit 745cc4ed6d
28 changed files with 33632 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mpq/cifar10_dws_cnn/cifar10_dws_cnn.py
+89
View File
@@ -0,0 +1,89 @@
import init_utils
import common
# Initialize the environment and get the name
name = init_utils.initialize_environment(__file__)
args = init_utils.get_args()
# Set arguments from command line
max_acc_drop = args.max_acc_drop
device = args.device
from sklearn.model_selection import train_test_split
import torch.nn as nn
import torch.nn.functional as F
import tensorflow as tf
import numpy as np
# Load our Dataset
(X_train, y_train), (X_test, y_test) = tf.keras.datasets.cifar10.load_data()
y_train = np.squeeze(y_train, axis = 1)
y_test = np.squeeze(y_test, axis = 1)
X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size = 0.15)
X_train = (np.transpose(X_train, (0,3,1,2)))
X_test = (np.transpose(X_test, (0,3,1,2)))
X_val = (np.transpose(X_val, (0,3,1,2)))
BATCH_SIZE = 128
epochs = 1
lr = 0.0001
class DepthwiseBlock(nn.Module):
def __init__(self, in_channels, out_channels):
super(DepthwiseBlock, self).__init__()
layers = []
layers.append(nn.Conv2d(in_channels = in_channels, out_channels = in_channels,
kernel_size = 3, padding = 1, groups = in_channels)) # Depthwise convolution
layers.append(nn.ReLU(inplace = True))
layers.append(nn.Conv2d(in_channels = in_channels, out_channels = out_channels,
kernel_size = 1, padding = 0)) # Pointwise convolution
layers.append(nn.ReLU(inplace = True))
self.block = nn.Sequential(*layers)
def forward(self, x):
return self.block(x)
class Cifar10_Dws_CNN(nn.Module):
def __init__(self):
super(Cifar10_Dws_CNN, self).__init__()
self.features = nn.Sequential(
DepthwiseBlock(in_channels = 3, out_channels = 64),
DepthwiseBlock(in_channels = 64, out_channels = 64),
nn.MaxPool2d(kernel_size = 2, stride = 2),
DepthwiseBlock(in_channels = 64, out_channels = 128),
DepthwiseBlock(in_channels = 128, out_channels = 128),
nn.MaxPool2d(kernel_size = 2, stride = 2),
DepthwiseBlock(in_channels = 128, out_channels = 256),
DepthwiseBlock(in_channels = 256, out_channels = 256),
nn.MaxPool2d(kernel_size = 2, stride = 2)
)
self.flatten = nn.Flatten()
self.classifier = nn.Sequential(
nn.Linear(256 * 4 * 4, 10) # Assuming input size is (32, 32) and after 3 max pooling layers, the size is (4, 4)
)
def forward(self, x):
x = self.features(x)
x = self.flatten(x)
x = self.classifier(x)
return F.log_softmax(x, dim = 1)
net = Cifar10_Dws_CNN()
common.create_ibex_qnn(net, name, device, X_train, y_train, X_test, y_test,
X_val = X_val, y_val = y_val, BATCH_SIZE = BATCH_SIZE,
epochs = epochs, lr = lr, max_acc_drop = max_acc_drop)
mpq/cmsis_cnn/cmsis_cnn.py
+77
View File
@@ -0,0 +1,77 @@
import init_utils
import common
# Initialize the environment and get the name
name = init_utils.initialize_environment(__file__)
args = init_utils.get_args()
# Set arguments from command line
max_acc_drop = args.max_acc_drop
device = args.device
from sklearn.model_selection import train_test_split
import torch.nn as nn
import torch.nn.functional as F
import tensorflow as tf
import numpy as np
# Load our Dataset
(X_train, y_train), (X_test, y_test) = tf.keras.datasets.cifar10.load_data()
y_train = np.squeeze(y_train, axis = 1)
y_test = np.squeeze(y_test, axis = 1)
X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size = 0.15)
X_train = (np.transpose(X_train, (0,3,1,2)) - 128.0)/255.0
X_test = (np.transpose(X_test, (0,3,1,2)) - 128.0)/255.0
X_val = (np.transpose(X_val, (0,3,1,2)) - 128.0)/255.0
BATCH_SIZE = 32
epochs = 1
lr = 0.0001
class CMSIS_CNN(nn.Module):
def __init__(self):
super(CMSIS_CNN, self).__init__()
self.conv1 = nn.Conv2d(in_channels = 3, out_channels = 32, kernel_size = 5, padding = 2)
self.relu1 = nn.ReLU()
self.max1 = nn.MaxPool2d(2,2)
self.d1 = nn.Dropout(p = 0.25)
self.conv2 = nn.Conv2d(in_channels = 32, out_channels = 32, kernel_size = 5, padding = 2)
self.relu2 = nn.ReLU()
self.max2 = nn.MaxPool2d(2,2)
self.d2 = nn.Dropout(p = 0.25)
self.conv3 = nn.Conv2d(in_channels = 32, out_channels = 64, kernel_size = 5, padding = 2)
self.relu3 = nn.ReLU()
self.max3 = nn.MaxPool2d(2,2)
self.d3 = nn.Dropout(p = 0.4)
self.flatten = nn.Flatten()
self.linear1 = nn.Linear(1024, 10)
def forward(self,X):
X = self.relu1((self.conv1(X)))
X = self.max1(X)
X = self.d1(X)
X = self.relu2((self.conv2(X)))
X = self.max2(X)
X = self.d2(X)
X = self.relu3((self.conv3(X)))
X = self.max3(X)
X = self.d3(X)
X = self.flatten(X)
X = self.linear1(X)
return F.log_softmax(X, dim = 1)
net = CMSIS_CNN()
common.create_ibex_qnn(net, name, device, X_train, y_train, X_test, y_test,
X_val = X_val, y_val = y_val, BATCH_SIZE = BATCH_SIZE,
epochs = epochs, lr = lr, max_acc_drop = max_acc_drop)
mpq/common.py
+10
View File
@@ -125,3 +125,13 @@ def create_ibex_qnn(net, name, device, X_train, y_train, X_test, y_test, X_val =
print('\nSIMULATING MODEL ON IBEX CORE\nUSE THE OUTPUTS TO VERIFY THAT THE RESULTS ARE CORRECT !!')
ibex_model = simulate_ibex.create_lenet_model(int_weights, int_og_bias, mul_vals, shift_vals)
simulate_ibex.eval_sim_model(quant_net, ibex_model, test_loader)
elif(name == 'cmsis_cnn'):
print('\nSIMULATING MODEL ON IBEX CORE\nUSE THE OUTPUTS TO VERIFY THAT THE RESULTS ARE CORRECT !!')
ibex_model = simulate_ibex.create_cmsis_cnn_model(int_weights, int_og_bias, mul_vals, shift_vals)
simulate_ibex.eval_sim_model(quant_net, ibex_model, test_loader)
elif(name == 'cifar10_dws_cnn'):
print('\nSIMULATING MODEL ON IBEX CORE\nUSE THE OUTPUTS TO VERIFY THAT THE RESULTS ARE CORRECT !!')
ibex_model = simulate_ibex.create_ibex_dws_model(int_weights, int_og_bias, mul_vals, shift_vals)
simulate_ibex.eval_sim_model(quant_net, ibex_model, test_loader)
mpq/configure_ibex.py
+171 -58
View File
@@ -27,26 +27,34 @@ def quantize_multiplier(real_multiplier):
return quantized_multiplier, right_shift
def get_int_params(quant_net):
int_weights = []
int_bias = []
in_scales = []
act_scales = []
for _, module in quant_net.sequential.named_children():
if hasattr(module, 'weight') and module.weight is not None:
int_weights.append(module.int_weight().cpu().numpy())
int_bias.append(module.int_bias().cpu().numpy())
in_scales.append(module.quant_bias_scale().cpu().detach().numpy())
def extract_quant_params(module):
for name, submodule in module.named_children():
# Check if the submodule has weights and append them if present
if hasattr(submodule, 'weight') and submodule.weight is not None:
int_weights.append(submodule.int_weight().cpu().detach().numpy())
int_bias.append(submodule.int_bias().cpu().detach().numpy())
in_scales.append(submodule.quant_bias_scale().cpu().detach().numpy())
if hasattr(module, 'quant_act_scale') and module.quant_act_scale is not None:
act_scales.append(module.quant_act_scale().cpu().detach().numpy())
# Check if the submodule has activation scale and append it if present
if hasattr(submodule, 'quant_act_scale') and submodule.quant_act_scale is not None:
act_scales.append(submodule.quant_act_scale().cpu().detach().numpy())
act_scales.append(quant_net.o_quant.quant_act_scale().cpu().detach().numpy())
# Recursively extract parameters from the children modules
extract_quant_params(submodule)
# Start extraction from the top-level module
extract_quant_params(quant_net)
mul_vals, shift_vals = [], []
for i in range(len(act_scales)):
M = in_scales[i]/act_scales[i]
for i in range(len(act_scales)-1):
M = in_scales[i]/act_scales[i+1]
mul, shift = quantize_multiplier(M[0])
mul_vals.append(mul)
shift_vals.append(shift)
@@ -87,22 +95,27 @@ def decide_mode(network, weight_bit_width, input_uint8 = True):
for name, module in network.named_modules():
if isinstance(module, layer_types_py):
layer_type_name = module.__class__.__name__
if(layer_type_name == 'Conv2d' or layer_type_name == 'Linear' or layer_type_name == 'DepthwiseConv2d'):
if(layer_type_name == 'Linear'):
layer_type.append(layer_type_name)
if(layer_type_name == 'Conv2d'):
if(module.groups == module.in_channels):
layer_type.append('DepthwiseConv2d')
else:
layer_type.append(layer_type_name)
else:
if(layer_type_name == 'ReLU' or layer_type_name == 'Sigmoid'):
input_sign[ins] = 0
ins += 1
for i in range(len(weight_bit_width)):
signed_input = 4 * input_sign[i]
if(weight_bit_width[i] == 2):
mode_per_layer.append(signed_input + 3)
elif(weight_bit_width[i] == 4):
mode_per_layer.append(signed_input + 2)
else:
if(layer_type[i] == 'DepthwiseConv2d'):
if(layer_type[i] == 'DepthwiseConv2d'):
mode_per_layer.append(signed_input + 1)
else:
if(weight_bit_width[i] == 2):
mode_per_layer.append(signed_input + 3)
elif(weight_bit_width[i] == 4):
mode_per_layer.append(signed_input + 2)
else:
mode_per_layer.append(signed_input)
@@ -161,15 +174,22 @@ def pad_inputs_weights(quant_net, test_loader, mode_per_layer,
else:
new_size_0 = a * 4
b = w.shape[1] // 4
if(w.shape[1] % 4 != 0):
new_size_1 = (b + 1) * 4
if((mode_per_layer[i] != 1) and (mode_per_layer[i] != 5)):
b = w.shape[1] // 4
if(w.shape[1] % 4 != 0):
new_size_1 = (b + 1) * 4
else:
new_size_1 = b * 4
new_w = np.zeros((new_size_0, new_size_1, w.shape[2], w.shape[3])).astype(np.int8)
new_w[:w.shape[0], :w.shape[1], :, :] = w
else:
new_size_1 = b * 4
new_w = np.zeros((new_size_0, new_size_1, w.shape[2], w.shape[3])).astype(np.int8)
new_w[:w.shape[0], :w.shape[1], :, :] = w
new_size_1 = 1
new_w = np.zeros((new_size_0, new_size_1, w.shape[2], w.shape[3])).astype(np.int8)
new_w[:w.shape[0], :w.shape[1], :, :] = w
new_w = np.squeeze(new_w, axis = 1)
padded_int_weights.append(new_w)
padded_int_biases = []
@@ -325,6 +345,15 @@ def concat_inputs_weights(mode_per_layer, padded_input, padded_int_weights, padd
comb = combine_values(vector)
new_mat[i][j] = comb
elif(len(dims) == 3):
new_mat = np.zeros((int(dims[0]//4), dims[1], dims[2]), dtype = np.int64)
for i in range(int(dims[0]//4)):
for j in range(dims[1]):
for k in range(dims[2]):
vector = layer_weight[4*i : 4*(i+1), j, k]
comb = combine_values(vector)
new_mat[i][j][k] = comb
elif(len(dims) == 4):
if((mode_per_layer[iter] == 0) | (mode_per_layer[iter] == 4)):
new_mat = np.zeros((int(dims[0]//4), dims[1], dims[2], dims[3]), dtype = np.int64)
@@ -602,9 +631,17 @@ def save_cnn_net_params(path, int_weights, int_biases, mul_vals, shift_vals, shi
dims = np.shape(int_weights[k])
mat = int_weights[k]
if(len(dims) == 2):
wi += 1
st = 'static const int W' + str(wi) + '[' + str(dims[0]) + ']' + '[' + str(dims[1]) + '] = {\n'
if(len(dims) == 2 or ((len(dims) == 4) and dims[2] == dims[3] == 1)):
f.write('static const int ')
if(len(dims) == 2):
wi += 1
f.write('W' + str(wi))
else:
mat = np.squeeze(mat, axis = (2,3))
fi += 1
f.write('F' + str(fi))
st = '[' + str(dims[0]) + ']' + '[' + str(dims[1]) + '] = {\n'
f.write(st)
for n in range(dims[0]):
f.write('\t{')
@@ -618,7 +655,33 @@ def save_cnn_net_params(path, int_weights, int_biases, mul_vals, shift_vals, shi
f.write(',')
f.write('\n')
f.write('};\n\n')
elif (len(dims) == 3):
dims = np.shape(mat)
fi += 1
st = 'static const int F' + str(fi) + '[' + str(dims[0]) + '][' + str(dims[1])
st += '][' + str(dims[2]) + '] = {\n'
f.write(st)
for n in range(dims[0]):
f.write('\t{\n')
for l in range(dims[1]):
f.write('\t\t{')
for h in range(dims[2] - 1):
f.write(str(mat[n][l][h]) + ', ')
if dims[2] != 1:
f.write(str(mat[n][l][dims[2] - 1]) + '}')
else:
f.write(str(mat[n][l][0]) + '}')
if (l != dims[1] - 1):
f.write(',')
f.write('\n')
f.write('\t}')
if n != dims[0] - 1:
f.write(',')
f.write('\n')
f.write('};\n\n')
elif(len(dims) == 4):
mat = np.transpose(mat, (0, 2, 3, 1))
dims = np.shape(mat)
@@ -856,9 +919,11 @@ def generate_opt_c_code_mlp(path, name, int_weights, optimal_config, type_of_lay
f.write('\t' + name + '();\n\n')
f.write('\treturn 0;\n}')
def get_cnn_details(model):
details = []
for layer in model.children():
def get_cnn_details(module, details = None):
if details is None:
details = []
for layer in module.children():
if isinstance(layer, nn.Conv2d):
details.append({
"layer_type": "Conv2d",
@@ -866,18 +931,19 @@ def get_cnn_details(model):
"out_channels": layer.out_channels,
"kernel_size": layer.kernel_size,
"stride": layer.stride,
"padding": layer.padding
"padding": layer.padding,
"groups": layer.groups
})
elif (isinstance(layer, nn.MaxPool2d)):
elif isinstance(layer, nn.MaxPool2d):
details.append({
"layer_type": "MaxPool2d",
"kernel_size": layer.kernel_size,
"stride": layer.stride,
"padding": layer.padding
})
elif (isinstance(layer, nn.AvgPool2d)):
elif isinstance(layer, nn.AvgPool2d):
details.append({
"layer_type": "AvgPool2d",
"kernel_size": layer.kernel_size,
@@ -891,6 +957,10 @@ def get_cnn_details(model):
"in_features": layer.in_features,
"out_features": layer.out_features
})
# Recursively apply to children modules
get_cnn_details(layer, details)
return details
def generate_og_c_code_cnn(path, name, input, cnn_details, int_weights):
@@ -900,10 +970,17 @@ def generate_og_c_code_cnn(path, name, input, cnn_details, int_weights):
f.write('#include "fully_connected.h"\n')
f.write('#include "ibex_cnn_params.h"\n')
f.write('#include "ibex_inputs.h"\n')
f.write('#include "conv2d.h"\n\n')
f.write('#include "conv2d.h"\n')
f.write('#define IMG_SZ ' + str(input.shape[2]) + '\n')
f.write('#define NUM_FIL0 ' + str(int_weights[0].shape[1]) + '\n\n')
for detail in cnn_details[:-1]:
if detail["layer_type"] == "Conv2d":
if(detail["in_channels"] == detail["out_channels"] == detail["groups"] != 1):
f.write('#include "dws_conv.h"\n')
break
f.write('\n')
f.write('#define IMG_SZ ' + str(np.shape(input)[2]) + '\n')
f.write('#define NUM_FIL0 ' + str(np.shape(input)[1]) + '\n\n')
i = 1
for w in int_weights:
if(len(np.shape(w)) == 4):
@@ -1050,11 +1127,17 @@ def generate_og_c_code_cnn(path, name, input, cnn_details, int_weights):
for detail in cnn_details[:-1]:
if detail["layer_type"] == "Conv2d":
if(detail["in_channels"] == detail["out_channels"] == detail["groups"] != 1):
conv_type = 'dw_conv'
elif(detail["kernel_size"][0] == 1):
conv_type = 'pw_conv'
else:
conv_type = "conv2"
if(i == 1):
f.write('\t\tconv2(inp_dim, f_dim1, outp_dim1, in, F1, B1, ')
f.write('\t\t' + conv_type + '(inp_dim, f_dim1, outp_dim1, in, F1, B1, ')
f.write('out1, STRIDE1, pad_1, SB1, MV1, SV1);')
else:
f.write('\t\tconv2(outp_dim' + str(i-1) + ', f_dim' + str(i) + ', outp_dim' + str(i))
f.write('\t\t' + conv_type + '(outp_dim' + str(i-1) + ', f_dim' + str(i) + ', outp_dim' + str(i))
f.write(', out' + str(i-1) + ', F' + str(fi) + ', B' + str(fi) + ', out' + str(i))
f.write(', STRIDE' + str(fi) + ', pad_' + str(i) + ', SB' + str(fi))
f.write(', MV' + str(fi) + ', SV' + str(fi) + ');')
@@ -1091,10 +1174,17 @@ def generate_og_c_code_cnn(path, name, input, cnn_details, int_weights):
f.write('\n')
i += 1
f.write('\t\tmlp_layer(out' + str(i-1) + ', out, DENSE_DIM' + str(dn-1))
f.write(', OUT_DIM, W' + str(dn) + ', B' + str(fi + dn - 1))
f.write(', SB' + str(fi + dn - 1) + ', MV' + str(fi + dn - 1))
f.write(', SV' + str(fi + dn - 1) + ');\n')
if flatten == 0:
f.write('\t\tflatten(outp_dim' + str(i-1) + ', out' + str(i-1) + ', out' + str(i) + ');\n\n')
i += 1
f.write('\t\tmlp_layer(out' + str(i-1) + ', out, flatten_dim, OUT_DIM, ')
f.write('W1, B' + str(fi + dn - 1) + ', SB' + str(fi + dn - 1) + ', MV' + str(fi + dn - 1))
f.write(', SV' + str(fi + dn - 1) + ');')
else:
f.write('\t\tmlp_layer(out' + str(i-1) + ', out, DENSE_DIM' + str(dn-1))
f.write(', OUT_DIM, W' + str(dn) + ', B' + str(fi + dn - 1))
f.write(', SB' + str(fi + dn - 1) + ', MV' + str(fi + dn - 1))
f.write(', SV' + str(fi + dn - 1) + ');\n')
f.write('\n\t\tpcount_enable(0);\n\n')
f.write('\t\tputs("Output Layer Values:\\n");\n')
@@ -1119,13 +1209,21 @@ def generate_opt_c_code_cnn(path, name, input, cnn_details, int_weights, optimal
f.write('#include "fully_connected_opt.h"\n')
f.write('#include "ibex_cnn_params.h"\n')
f.write('#include "ibex_inputs.h"\n')
f.write('#include "conv2d_opt.h"\n\n')
f.write('#include "conv2d_opt.h"\n')
for detail in cnn_details[:-1]:
if detail["layer_type"] == "Conv2d":
if(detail["in_channels"] == detail["out_channels"] == detail["groups"] != 1):
f.write('#include "dws_conv_opt.h"\n')
break
f.write('\n')
f.write('#define IMG_SZ ' + str(np.shape(input)[2]) + '\n')
f.write('#define NUM_FIL0 ' + str(np.shape(input)[0]) + '\n\n')
f.write('#define NUM_FIL0 ' + str(np.shape(input)[1]) + '\n\n')
i = 1
for w in int_weights:
if(len(np.shape(w)) == 4):
if(len(np.shape(w)) == 4 or len(np.shape(w)) == 3):
f.write('#define FILTER' + str(i) + ' ' + str(w.shape[2]) + '\n')
i += 1
@@ -1133,7 +1231,7 @@ def generate_opt_c_code_cnn(path, name, input, cnn_details, int_weights, optimal
i = 1
for w in int_weights:
if(len(np.shape(w)) == 4):
if(len(np.shape(w)) == 4 or len(np.shape(w)) == 3):
f.write('#define NUM_FIL' + str(i) + ' ' + str(w.shape[0]) + '\n')
i += 1
@@ -1270,14 +1368,21 @@ def generate_opt_c_code_cnn(path, name, input, cnn_details, int_weights, optimal
for detail in cnn_details[:-1]:
if detail["layer_type"] == "Conv2d":
if(detail["in_channels"] == detail["out_channels"] == detail["groups"] != 1):
conv_type = 'dw_conv_opt'
elif(detail["kernel_size"][0] == 1):
conv_type = 'pw_conv_' + str(optimal_config[j]) + 'bits'
else:
conv_type = 'conv2_' + str(optimal_config[j]) + 'bits'
if(i == 1):
f.write('\t\tconv2_' + str(optimal_config[j]) + 'bits')
if(np.shape(input)[0] == 1):
f.write('\t\t' + conv_type)
if(np.shape(input)[1] == 1):
f.write('_1ch')
f.write('(inp_dim, f_dim1, outp_dim1, in, F1, B1, ')
f.write('out1, STRIDE1, pad_1, SB1, MV1, SV1);')
else:
f.write('\t\tconv2_' + str(optimal_config[j]) + 'bits(outp_dim' + str(i-1) + ', f_dim' + str(i))
f.write('\t\t' + conv_type + '(outp_dim' + str(i-1) + ', f_dim' + str(i))
f.write(', outp_dim' + str(i) + ', out' + str(i-1) + ', F' + str(fi) + ', B' + str(fi) + ', out')
f.write(str(i) + ', STRIDE' + str(fi) + ', pad_' + str(i) + ', SB' + str(fi))
f.write(', MV' + str(fi) + ', SV' + str(fi) + ');')
@@ -1314,11 +1419,19 @@ def generate_opt_c_code_cnn(path, name, input, cnn_details, int_weights, optimal
f.write('\n')
i += 1
f.write('\t\tmlp_layer_' + str(optimal_config[-1]) + 'bits(out' + str(i-1) + ', out, DENSE_DIM' + str(dn-1))
f.write(', OUT_DIM, W' + str(dn) + ', B' + str(fi + dn - 1))
f.write(', SB' + str(fi + dn - 1) + ', MV' + str(fi + dn - 1))
f.write(', SV' + str(fi + dn - 1) + ');\n')
if flatten == 0:
f.write('\t\tflatten(outp_dim' + str(i-1) + ', out' + str(i-1) + ', out' + str(i) + ');\n\n')
i += 1
f.write('\t\tmlp_layer_' + str(optimal_config[j]) + 'bits(out' + str(i-1) + ', out, ')
f.write('flatten_dim, OUT_DIM, W1, B' + str(fi + dn - 1) + ', SB' + str(fi + dn - 1) + ', MV')
f.write(str(fi + dn - 1) + ', SV' + str(fi + dn - 1) + ');\n')
else:
f.write('\t\tmlp_layer_' + str(optimal_config[-1]) + 'bits(out' + str(i-1) + ', out, DENSE_DIM' + str(dn-1))
f.write(', OUT_DIM, W' + str(dn) + ', B' + str(fi + dn - 1))
f.write(', SB' + str(fi + dn - 1) + ', MV' + str(fi + dn - 1))
f.write(', SV' + str(fi + dn - 1) + ');\n')
f.write('\n\t\tpcount_enable(0);\n\n')
f.write('\t\tputs("Output Layer Values:\\n");\n')
f.write('\t\tfor(int i = 0; i < OUT_DIM; i++) {\n')
f.write('\t\t\tputhex((out[i] & 0xFF000000) >> 24);\n')
mpq/mpq_quantize.py
+113 -22
View File
@@ -11,6 +11,9 @@ from torch import nn, optim
import brevitas.nn as qnn
from brevitas.quant import *
from brevitas.core.restrict_val import RestrictValueType
from collections import defaultdict
from torchinfo import summary
def net_input_size(X_train):
@@ -202,7 +205,21 @@ def generate_sequences(length, values = [2, 4, 8]):
def create_weight_confs(macc_per_layer):
total_macc_opt = []
weights_per_layer = generate_sequences(len(macc_per_layer))
cc = 0
idx = []
if(len(macc_per_layer) >= 6):
for i, mpl in enumerate(macc_per_layer):
if(mpl/max(macc_per_layer) < 0.05):
cc += 1
idx.append(i)
weights_per_layer = generate_sequences(len(macc_per_layer) - cc)
for w in weights_per_layer:
for i in idx:
w.insert(i, 8)
for w_conf in weights_per_layer:
macc = 0
@@ -234,24 +251,47 @@ def create_weight_confs(macc_per_layer):
# Define a mapping from PyTorch layers to Brevitas layers
def create_layer_mapping(bit_width):
mapping = {
nn.Conv2d: lambda layer, bw: qnn.QuantConv2d(in_channels = layer.in_channels,
out_channels = layer.out_channels,
kernel_size = layer.kernel_size,
stride = layer.stride[0],
padding = layer.padding,
bias = True,
cache_inference_bias = True,
bias_quant = Int32Bias,
weight_bit_width = bw,
weight_quant = Int8WeightPerTensorFloat),
nn.Conv2d: lambda layer, bw: (qnn.QuantConv2d(in_channels=layer.in_channels,
out_channels=layer.out_channels,
kernel_size=layer.kernel_size,
stride=layer.stride[0],
padding=layer.padding,
groups=layer.groups,
bias=True,
cache_inference_bias=True,
bias_quant=Int32Bias,
weight_bit_width=bw,
weight_quant=Int8WeightPerTensorFloat,
weight_scaling_min_val=2e-16,
restrict_scaling_type=RestrictValueType.LOG_FP,
return_quant_tensor=True
) if layer.groups != layer.in_channels else (
# Special case for depthwise convolutions
qnn.QuantConv2d(in_channels=layer.in_channels,
out_channels=layer.out_channels,
kernel_size=layer.kernel_size,
stride=layer.stride[0],
padding=layer.padding,
groups=layer.groups,
bias=True,
cache_inference_bias=True,
bias_quant=Int32Bias,
weight_bit_width=8, # Fixed bit width for depthwise convolutions
weight_quant=Int8WeightPerTensorFloat,
weight_scaling_min_val=2e-16,
restrict_scaling_type=RestrictValueType.LOG_FP,
return_quant_tensor=True))),
nn.Linear: lambda layer, bw: qnn.QuantLinear(in_features = layer.in_features,
out_features = layer.out_features,
cache_inference_bias = True,
weight_quant = Int8WeightPerTensorFloat,
cache_inference_bias = True,
bias_quant = Int32Bias,
bias = True,
weight_bit_width = bw),
weight_quant = Int8WeightPerTensorFloat,
weight_bit_width = bw,
return_quant_tensor=True),
nn.ReLU: lambda _, bw: qnn.QuantReLU(bit_width = bw,
return_quant_tensor = True),
@@ -278,13 +318,11 @@ def convert_layer(layer, bit_width, layer_mapping):
return layer
# Function to convert a PyTorch model to a Brevitas model
def convert_model(module, bit_widths, layer_mapping):
layer_idx = [0]
def convert_model(module, bit_widths, layer_mapping, layer_idx = [0]):
brevitas_module = nn.Sequential()
for name, layer in module.named_children():
if list(layer.children()): # If the layer has children, recurse
brevitas_module.add_module(name, convert_model(layer, bit_widths, layer_mapping))
brevitas_module.add_module(name, convert_model(layer, bit_widths, layer_mapping, layer_idx))
else:
layer_type = type(layer)
if layer_type in [nn.Conv2d, nn.Linear]:
@@ -293,6 +331,7 @@ def convert_model(module, bit_widths, layer_mapping):
else:
bit_width = 8
brevitas_module.add_module(name, convert_layer(layer, bit_width, layer_mapping))
return brevitas_module
class Quant_Model(nn.Module):
@@ -300,13 +339,15 @@ class Quant_Model(nn.Module):
super(Quant_Model, self).__init__()
if(input_sign):
self.quant_inp = qnn.QuantIdentity(bit_width = 8, return_quant_tensor = True,
act_quant = Uint8ActPerTensorFloat)
act_quant = Uint8ActPerTensorFloat, scaling_min_val = 2e-16,
restrict_scaling_type = RestrictValueType.LOG_FP)
else:
self.quant_inp = qnn.QuantIdentity(bit_width = 8, return_quant_tensor = True,
act_quant = Int8ActPerTensorFloat)
act_quant = Int8ActPerTensorFloat, scaling_min_val = 2e-16,
restrict_scaling_type = RestrictValueType.LOG_FP)
self.sequential = convert_model(og_model, w, layer_mapping)
self.sequential = convert_model(og_model, w, layer_mapping, [0])
self.o_quant = qnn.QuantIdentity(bit_width = 8, return_quant_tensor = True)
def forward(self, X):
@@ -314,7 +355,36 @@ class Quant_Model(nn.Module):
X = self.sequential(X)
X = self.o_quant(X)
return F.log_softmax(X, dim = 1)
def count_layers_in_sequential(module):
# List to store the counts of Conv2d and Linear layers for each nn.Sequential module
sequential_counts = []
def _count_layers(submodule, prefix = ''):
if isinstance(submodule, nn.Sequential):
conv_count = 0
linear_count = 0
# Count Conv2d and Linear layers in the current nn.Sequential module
for child in submodule.children():
if isinstance(child, nn.Conv2d):
conv_count += 1
elif isinstance(child, nn.Linear):
linear_count += 1
# Append the counts to the list
sequential_counts.append((conv_count, linear_count))
# Recursively process children of the current nn.Sequential module
for name, child in submodule.named_children():
child_prefix = f"{prefix}.{name}" if prefix else name
_count_layers(child, child_prefix)
else:
# Process children of non-nn.Sequential modules
for name, child in submodule.named_children():
_count_layers(child, prefix)
_count_layers(module)
return sequential_counts[1:]
def train_quant_model(quant_net, train_loader, val_loader = None, device = 'cpu',
epochs = 20, lr = 0.0001):
@@ -392,6 +462,7 @@ def dse(og_model, max_acc_drop, weights_per_layer, fp_accuracy, train_loader, te
device = 'cpu', epochs = 5, lr = 0.0001):
sign = calculate_minimum(train_loader) >= 0
seq_counts = count_layers_in_sequential(og_model)
if max_acc_drop is not None:
print('\nDSE STARTING ... BINARY SEARCH')
@@ -402,6 +473,16 @@ def dse(og_model, max_acc_drop, weights_per_layer, fp_accuracy, train_loader, te
mid = (low + high) // 2
w = weights_per_layer[mid]
f_w = []
for i in range(len(seq_counts)):
t_w = w[i]
c,l = seq_counts[i]
for j in range(c+l):
f_w.append(t_w)
if(len(seq_counts) > 0):
w = f_w
# Create and train the quantized network
layer_mapping = create_layer_mapping(w)
quant_net = Quant_Model(og_model, w, layer_mapping, sign)
@@ -436,6 +517,16 @@ def dse(og_model, max_acc_drop, weights_per_layer, fp_accuracy, train_loader, te
print('\nDSE STARTING ... EXHAUSTIVE SEARCH')
test_accuracy = []
for i, w in enumerate(weights_per_layer):
f_w = []
for i in range(len(seq_counts)):
t_w = w[i]
c,l = seq_counts[i]
for j in range(c+l):
f_w.append(t_w)
if(len(seq_counts) > 0):
w = f_w
layer_mapping = create_layer_mapping(w)
quant_net = Quant_Model(og_model, w, layer_mapping, sign)
quant_net = quant_net.to(device)
mpq/simulate_ibex.py
+166 -26
View File
@@ -151,15 +151,154 @@ class Ibex_Lenet5(nn.Module):
return X
class Ibex_CMSIS_CNN(nn.Module):
def __init__(self, mul_vals, shift_vals):
super(Ibex_CMSIS_CNN, self).__init__()
self.m0 = mul_vals[0]
self.m1 = mul_vals[1]
self.m2 = mul_vals[2]
self.m3 = mul_vals[3]
self.s0 = shift_vals[0] + 7
self.s1 = shift_vals[1] + 7
self.s2 = shift_vals[2] + 7
self.s3 = shift_vals[3] + 7
self.conv1 = nn.Conv2d(in_channels = 3, out_channels = 32, kernel_size = 5, padding = 'same')
self.max1 = nn.MaxPool2d(2,2)
self.conv2 = nn.Conv2d(in_channels = 32, out_channels = 32, kernel_size = 5, padding = 'same')
self.max2 = nn.MaxPool2d(2,2)
self.conv3 = nn.Conv2d(in_channels = 32, out_channels = 64, kernel_size = 5, padding = 'same')
self.max3 = nn.MaxPool2d(2,2)
self.linear1 = nn.Linear(1024, 10)
def forward(self, X, print_out = False):
X = self.conv1(X)
X = torch.mul(X, self.m0)
X = torch.add(X, torch.bitwise_left_shift(torch.tensor(1), self.s0 -1)).type(torch.LongTensor)
X = torch.bitwise_right_shift(X, self.s0).type(torch.FloatTensor)
X = torch.clamp(X, min = 0, max = 255)
X = self.max1(X)
X = self.conv2(X)
X = torch.mul(X, self.m1)
X = torch.add(X, torch.bitwise_left_shift(torch.tensor(1), self.s1 -1)).type(torch.LongTensor)
X = torch.bitwise_right_shift(X, self.s1).type(torch.FloatTensor)
X = torch.clamp(X, min = 0, max = 255)
X = self.max2(X)
X = self.conv3(X)
X = torch.mul(X, self.m2)
X = torch.add(X, torch.bitwise_left_shift(torch.tensor(1), self.s2 -1)).type(torch.LongTensor)
X = torch.bitwise_right_shift(X, self.s2).type(torch.FloatTensor)
X = torch.clamp(X, min = 0, max = 255)
X = self.max3(X)
X = X.reshape(X.shape[0], -1)
X = self.linear1(X)
X = torch.mul(X, self.m3)
X = torch.add(X, torch.bitwise_left_shift(torch.tensor(1), self.s3 -1)).type(torch.LongTensor)
X = torch.bitwise_right_shift(X, self.s3).type(torch.FloatTensor)
X = torch.clamp(X, min = 0, max = 255)
if(print_out):
print(X)
return X
class Ibex_DepthwiseBlock(nn.Module):
def __init__(self, in_channels, out_channels, mul_vals, shift_vals):
super(Ibex_DepthwiseBlock, self).__init__()
self.dw = nn.Conv2d(in_channels = in_channels, out_channels = in_channels,
kernel_size = 3, padding = 1, groups = in_channels)
self.pw = nn.Conv2d(in_channels = in_channels, out_channels = out_channels,
kernel_size = 1, padding = 0)
self.m0 = mul_vals[0]
self.m1 = mul_vals[1]
self.s0 = shift_vals[0] + 7
self.s1 = shift_vals[1] + 7
def forward(self, X):
X = self.dw(X)
X = torch.mul(X, self.m0)
X = torch.add(X, torch.bitwise_left_shift(torch.tensor(1), self.s0 -1)).type(torch.LongTensor)
X = torch.bitwise_right_shift(X, self.s0).type(torch.FloatTensor)
X = torch.clamp(X, min = 0, max = 255)
X = self.pw(X)
X = torch.mul(X, self.m1)
X = torch.add(X, torch.bitwise_left_shift(torch.tensor(1), self.s1 -1)).type(torch.LongTensor)
X = torch.bitwise_right_shift(X, self.s1).type(torch.FloatTensor)
X = torch.clamp(X, min = 0, max = 255)
return X
class Ibex_Cifar10_Dws_CNN(nn.Module):
def __init__(self, mul_vals, shift_vals):
super(Ibex_Cifar10_Dws_CNN, self).__init__()
self.features = nn.Sequential(
Ibex_DepthwiseBlock(3, 64, mul_vals[0:2], shift_vals[0:2]),
Ibex_DepthwiseBlock(64, 64, mul_vals[2:4], shift_vals[2:4]),
nn.MaxPool2d(kernel_size = 2, stride = 2),
Ibex_DepthwiseBlock(64, 128, mul_vals[4:6], shift_vals[4:6]),
Ibex_DepthwiseBlock(128, 128, mul_vals[6:8], shift_vals[6:8]),
nn.MaxPool2d(kernel_size = 2, stride = 2),
Ibex_DepthwiseBlock(128, 256, mul_vals[8:10], shift_vals[8:10]),
Ibex_DepthwiseBlock(256, 256, mul_vals[10:12], shift_vals[10:12]),
nn.MaxPool2d(kernel_size = 2, stride = 2)
)
self.flatten = nn.Flatten()
self.classifier = nn.Sequential(
nn.Linear(256 * 4 * 4, 10)
)
self.m_cl = mul_vals[12]
self.s_cl = shift_vals[12] + 7
def forward(self, x, print_out = False):
x = self.features(x)
x = self.flatten(x)
x = self.classifier(x)
x = torch.mul(x, self.m_cl)
x = torch.add(x, torch.bitwise_left_shift(torch.tensor(1), self.s_cl - 1)).type(torch.LongTensor)
x = torch.bitwise_right_shift(x, self.s_cl).type(torch.FloatTensor)
x = torch.clamp(x, min = 0, max = 255)
if(print_out):
print(x)
return x
def configure_network(ibex_model_dict, int_weights, int_biases):
for i, (name, _) in enumerate(ibex_model_dict.items()):
if(i%2 == 0):
ibex_model_dict[name] = torch.tensor(int_weights[i//2])
else:
ibex_model_dict[name] = torch.tensor(int_biases[i//2])
return ibex_model_dict
def create_fann_model(int_weights, int_biases, mul_vals, shift_vals):
ibex_model = Ibex_FANN(mul_vals, shift_vals)
ibex_model_dict = ibex_model.state_dict()
ibex_model_dict['linear1.weight'] = torch.tensor(int_weights[0])
ibex_model_dict['linear2.weight'] = torch.tensor(int_weights[1])
ibex_model_dict['linear1.bias'] = torch.tensor(int_biases[0])
ibex_model_dict['linear2.bias'] = torch.tensor(int_biases[1])
ibex_model_dict = configure_network(ibex_model_dict, int_weights, int_biases)
ibex_model.load_state_dict(ibex_model_dict)
return ibex_model
@@ -168,15 +307,7 @@ def create_uci_model(int_weights, int_biases, mul_vals, shift_vals):
ibex_model = Ibex_UCI_MLP(mul_vals, shift_vals)
ibex_model_dict = ibex_model.state_dict()
ibex_model_dict['fc0.weight'] = torch.tensor(int_weights[0])
ibex_model_dict['fc1.weight'] = torch.tensor(int_weights[1])
ibex_model_dict['fc2.weight'] = torch.tensor(int_weights[2])
ibex_model_dict['fc3.weight'] = torch.tensor(int_weights[3])
ibex_model_dict['fc0.bias'] = torch.tensor(int_biases[0])
ibex_model_dict['fc1.bias'] = torch.tensor(int_biases[1])
ibex_model_dict['fc2.bias'] = torch.tensor(int_biases[2])
ibex_model_dict['fc3.bias'] = torch.tensor(int_biases[3])
ibex_model_dict = configure_network(ibex_model_dict, int_weights, int_biases)
ibex_model.load_state_dict(ibex_model_dict)
@@ -186,24 +317,33 @@ def create_lenet_model(int_weights, int_biases, mul_vals, shift_vals):
ibex_model = Ibex_Lenet5(mul_vals, shift_vals)
ibex_model_dict = ibex_model.state_dict()
ibex_model_dict['conv1.weight'] = torch.tensor(int_weights[0])
ibex_model_dict['conv2.weight'] = torch.tensor(int_weights[1])
ibex_model_dict['fc1.weight'] = torch.tensor(int_weights[2])
ibex_model_dict['fc2.weight'] = torch.tensor(int_weights[3])
ibex_model_dict['fc3.weight'] = torch.tensor(int_weights[4])
ibex_model_dict['conv1.bias'] = torch.tensor(int_biases[0])
ibex_model_dict['conv2.bias'] = torch.tensor(int_biases[1])
ibex_model_dict['fc1.bias'] = torch.tensor(int_biases[2])
ibex_model_dict['fc2.bias'] = torch.tensor(int_biases[3])
ibex_model_dict['fc3.bias'] = torch.tensor(int_biases[4])
ibex_model_dict = configure_network(ibex_model_dict, int_weights, int_biases)
ibex_model.load_state_dict(ibex_model_dict)
return ibex_model
def create_cmsis_cnn_model(int_weights, int_biases, mul_vals, shift_vals):
ibex_model = Ibex_CMSIS_CNN(mul_vals, shift_vals)
ibex_model_dict = ibex_model.state_dict()
ibex_model_dict = configure_network(ibex_model_dict, int_weights, int_biases)
ibex_model.load_state_dict(ibex_model_dict)
return ibex_model
def create_ibex_dws_model(int_weights, int_biases, mul_vals, shift_vals):
ibex_model = Ibex_Cifar10_Dws_CNN(mul_vals, shift_vals)
ibex_model_dict = ibex_model.state_dict()
ibex_model_dict = configure_network(ibex_model_dict, int_weights, int_biases)
ibex_model.load_state_dict(ibex_model_dict)
return ibex_model
def eval_sim_model(quant_model, ibex_model, test_loader):
# Turn off gradients for validation
with torch.no_grad():
ibex_model.eval()
correct = 0