def quantize_layer(q_fp_layer, wide, activation, num_bits):
'''
Quantizes layer
'''
ai8x.set_device(device=85, simulate=True, round_avg=False, verbose=False)
in_channels = q_fp_layer.op.weight.shape[1]
out_channels = q_fp_layer.op.weight.shape[0]
kernel_size = q_fp_layer.op.weight.shape[2:]
q_int_layer = ai8x.Conv2d(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
bias=False,
wide=wide,
activation=activation,
weight_bits=num_bits,
bias_bits=8,
quantize_activation=True)
out_shift = q_fp_layer.calc_out_shift(q_fp_layer.op.weight.detach(),
q_fp_layer.output_shift.detach())
weight_scale = q_fp_layer.calc_weight_scale(out_shift)
ai8x.set_device(device=85, simulate=False, round_avg=False, verbose=False)
weight = q_fp_layer.clamp_weight(q_fp_layer.quantize_weight(weight_scale *
q_fp_layer.op.weight))
q_int_weight = (2**(num_bits-1)) * weight
q_int_layer.output_shift = torch.nn.Parameter(
-torch.log2(weight_scale) # pylint: disable=no-member
)
q_int_layer.op.weight = torch.nn.Parameter(q_int_weight)
return q_int_layer
def test():
'''
Main test function
'''
wide_opts = [False, True]
act_opts = [None, 'ReLU']
bit_opts = [8, 4, 2, 1]
inp, inp_int = create_input_data(512)
for bit in bit_opts:
for act in act_opts:
for wide in wide_opts:
if wide and (act is not None):
continue
print(f'Testing for bits:{bit}, wide:{wide}, activation:{act} ...', end=' ')
fp_layer = create_conv2d_layer(512, 16, 3, wide, act)
q_fp_layer = quantize_fp_layer(fp_layer, wide, act, bit)
q_int_layer = quantize_layer(q_fp_layer, wide, act, bit)
ai8x.set_device(device=85, simulate=False, round_avg=False, verbose=False)
q_fp_out = q_fp_layer(inp)
ai8x.set_device(device=85, simulate=True, round_avg=False, verbose=False)
q_int_out = q_int_layer(inp_int)
assert ((128. * q_fp_out) == q_int_out).all(), 'FAIL!!'
print('PASS')
print('\nSUCCESS!!')
def test():
"""Test routine for nas implementation"""
ai8x.set_device(device=85, simulate=False, round_avg=False, verbose=False)
test_unit2d()
test_unit1d()
test_elastic_kernel_2d()
test_elastic_kernel_1d()
test_elastic_depth_2d()
test_elastic_depth_1d()
test_elastic_width_2d()
def main():
"""Main routine"""
ai8x.set_device(device=85, simulate=False, round_avg=False, verbose=False)
supported_models, model_names = load_models()
supported_sources, dataset_names = load_datasets()
args = parse_args(model_names, dataset_names)
args.truncate_testset = False
use_cuda = torch.cuda.is_available()
args.device = torch.device("cuda:0" if use_cuda else "cpu")
args.act_mode_8bit = False
# Get policy for once for all training policy
nas_policy = parse_nas_yaml.parse(args.nas_policy) \
if args.nas_policy.lower() != '' else None
# Get data loaders
train_loader, val_loader = get_data_loaders(supported_sources, args)
# Load model
model = create_model(supported_models, args)
checkpoint = torch.load(args.model_path, map_location=args.device)
model.load_state_dict(checkpoint['state_dict'])
# Calculate full model accuracy
full_model_acc = nas_utils.calc_accuracy(None, model, train_loader,
val_loader, args.device)
print(f'Model Accuracy: {100*full_model_acc: .3f}%')
# Run evolutionary search to find proper networks
evo_search_params = get_evo_search_params(nas_policy)
evo_search = EvolutionSearch(
population_size=evo_search_params['population_size'],
prob_mutation=evo_search_params['prob_mutation'],
ratio_mutation=evo_search_params['ratio_mutation'],
ratio_parent=evo_search_params['ratio_parent'],
num_iter=evo_search_params['num_iter'])
evo_search.set_model(model)
arch_list = evo_search.run(evo_search_params['constraints'], train_loader,
val_loader, args.device)
if args.export_archs:
generate_out_file(arch_list, min(args.num_out_archs, len(arch_list)),
args.dimensions, args.model_type, args.arch_file)
else:
for idx in range(min(args.num_out_archs, len(arch_list))):
print(f'Model-{idx}:')
print(f'\tArch: {arch_list[idx][0]}')
print(f'\tAcc: {arch_list[idx][1]}\n')
def quantize_fp_layer(fp_layer, wide, activation, num_bits):
'''
Creates layer with quantized leveled fp32 weights from a fp32 weighted layer
'''
ai8x.set_device(device=85, simulate=False, round_avg=False, verbose=False)
in_channels = fp_layer.op.weight.shape[1]
out_channels = fp_layer.op.weight.shape[0]
kernel_size = fp_layer.op.weight.shape[2:]
q_fp_layer = ai8x.Conv2d(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
bias=False,
wide=wide,
activation=activation,
weight_bits=num_bits,
bias_bits=8,
quantize_activation=True)
q_fp_layer.op.weight = copy.deepcopy(fp_layer.op.weight)
return q_fp_layer
def create_conv2d_layer(in_channels, out_channels, kernel_size, wide, activation):
'''
Creates randomly initialized layer
'''
ai8x.set_device(device=85, simulate=False, round_avg=False, verbose=False)
fp_layer = ai8x.Conv2d(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
bias=False,
wide=wide,
activation=activation)
fp_layer.op.weight = torch.nn.Parameter(
(2.0 * torch.rand(out_channels, # pylint: disable=no-member
in_channels,
kernel_size,
kernel_size) - 1.0)
)
return fp_layer
