KeyError when executing quantization

[iksooman]

Hi!
I'm trying to execute quick_start_for_expert.py of example with MobileNetV3.
But I got an Error like the below during quantization:

Validation Acc@1 86.400
Test: [0/40]    Time 0.50012    Acc@1 88.67188
Test: [10/40]   Time 0.05412    Acc@1 86.78977
Test: [20/40]   Time 0.03281    Acc@1 86.66295
Test: [30/40]   Time 0.02411    Acc@1 86.70615
Validation Acc@1 86.710
Start preparing the model for quantization
Traceback (most recent call last):
  File "custom_start.py", line 149, in <module>
    main_worker(args)
  File "custom_start.py", line 101, in main_worker
    qat_model = quantizer.quantize()
  File "/home/sss/embed_deep_lec/TinyNeuralNetwork/examples/../tinynn/graph/quantization/quantizer.py", line 487, in quantize
    self.rewrite_quantize_graph(graph)
  File "/home/sss/embed_deep_lec/TinyNeuralNetwork/examples/../tinynn/graph/quantization/quantizer.py", line 2879, in rewrite_quantize_graph
    node = graph.nodes_map[name]
KeyError: 'featureList_8_conv1_2_avg_pool'

Could you give me an idea for fixing it?

[iksooman]

Is it OK to modify the line 2879 of TinyNeuralNetwork/tinynn/graph/quantization/quantizer.py?

from

node = graph.nodes_map[name]

to

try:
    node = graph.nodes_map[name]
except KeyError:
    continue

[peterjc123]

@iksooman No, I think the correct thing to do is to remove all the things in the specific directory. since it usually mean that the model is changed or updated. I think the following logic would be better.

if name in graph.nodes_map:
    node = graph.nodes_map[name]
else:
    log.error(f'Node name {name} not found in configuration file, it probably means that your model has been updated. Please remove the old yaml file and try again')
    assert False

[iksooman]

@peterjc123 Could you please give me a little more detailed explanation?

I just replaced the model in the quick_start_for_expert.py file. Between the part where the model is defined and the execution of quantizer.quantize(), there seems to be only pruning. Do I need to modify the settings in pruning part?

def main_worker(args):
    print("###### TinyNeuralNetwork quick start for expert ######")

    # If you encounter any problems, please set the global log level to `DEBUG`, which may make it easier to debug.
    # set_global_log_level("DEBUG")

    model = mobilenet.Mobilenet()
    model.load_state_dict(torch.load(mobilenet.DEFAULT_STATE_DICT))

    device = get_device()
    model.to(device=device)

    if args.distillation:
        teacher = copy.deepcopy(model)

    if args.parallel:
        model = nn.DataParallel(model)

    # Provide a viable input for the model
    dummy_input = torch.rand((1, 3, 224, 224))

    context = DLContext()
    context.device = device
    context.train_loader, context.val_loader = get_dataloader(args.data_path, 224, args.batch_size, args.workers)

    print("Validation accuracy of the original model")
    validate(model, context)

    print("Start pruning the model")
    # If you need to set the sparsity of a single operator, then you may refer to the examples in `examples/pruner`.
    pruner = OneShotChannelPruner(model, dummy_input, {"sparsity": 0.75, "metrics": "l2_norm"})

    st_flops = pruner.calc_flops()
    pruner.prune()  # Get the pruned model

    print("Validation accuracy of the pruned model")
    validate(model, context)

    ed_flops = pruner.calc_flops()
    print(f"Pruning over, reduced FLOPS {100 * (st_flops - ed_flops) / st_flops:.2f}%  ({st_flops} -> {ed_flops})")

    print("Start finetune the pruned model")
    # In our experiments, using the same learning rate configuration as the one used during training from scratch
    # leads to a higher final model accuracy.
    context.max_epoch = 220
    context.criterion = nn.BCEWithLogitsLoss()
    context.optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
    context.scheduler = CosineAnnealingLR(context.optimizer, T_max=context.max_epoch + 1, eta_min=0)

    if args.warmup:
        print("Use warmup")
        context.warmup_iteration = len(context.train_loader) * 10  # warmup 10 epoch
        context.warmup_scheduler = CyclicLR(
            context.optimizer, base_lr=0, max_lr=0.1, step_size_up=context.warmup_iteration
        )

    if args.distillation:
        # The utilization of distillation may leads to better accuracy at the price of longer training time.
        print("Use distillation")
        context.custom_args = {'distill_A': 0.3, 'distill_T': 6, 'distill_teacher': teacher}
        train(model, context, train_one_epoch_distill, validate)
    else:
        train(model, context, train_one_epoch, validate)

    print("Start preparing the model for quantization")
    # We provides a QATQuantizer class that may rewrite the graph for and perform model fusion for quantization
    # The model returned by the `quantize` function is ready for QAT training
    quantizer = QATQuantizer(model, dummy_input, work_dir='out')
    qat_model = quantizer.quantize()

[peterjc123]

@iksooman you just have to do rm -rf out before running the code for a new model.

[iksooman]

@peterjc123 problem solved. Thank you!