@pytest.mark.skipif(os.getenv('NM_ML_SKIP_PYTORCH_TESTS', False), reason='Skipping pytorch tests')
@pytest.mark.skipif(os.getenv('NM_ML_SKIP_PYTORCH_QUANT_TESTS', False), reason='Skipping pytorch torch quantization tests')
@pytest.mark.skipif((torch_quantization is None), reason='torch quantization not available')
def test_quantization_modifier_yaml():
    start_epoch = 0.0
    submodules = ['block.0', 'block.2']
    model_fuse_fn_name = 'fuse_module'
    disable_quantization_observer_epoch = 2.0
    freeze_bn_stats_epoch = 3.0
    yaml_str = '\n        !QuantizationModifier\n            start_epoch: {start_epoch}\n            submodules: {submodules}\n            model_fuse_fn_name: {model_fuse_fn_name}\n            disable_quantization_observer_epoch: {disable_quantization_observer_epoch}\n            freeze_bn_stats_epoch: {freeze_bn_stats_epoch}\n        '.format(start_epoch=start_epoch, submodules=submodules, model_fuse_fn_name=model_fuse_fn_name, disable_quantization_observer_epoch=disable_quantization_observer_epoch, freeze_bn_stats_epoch=freeze_bn_stats_epoch)
    yaml_modifier = QuantizationModifier.load_obj(yaml_str)
    serialized_modifier = QuantizationModifier.load_obj(str(yaml_modifier))
    obj_modifier = QuantizationModifier(start_epoch=start_epoch, submodules=submodules, model_fuse_fn_name=model_fuse_fn_name, disable_quantization_observer_epoch=disable_quantization_observer_epoch, freeze_bn_stats_epoch=freeze_bn_stats_epoch)
    assert isinstance(yaml_modifier, QuantizationModifier)
    assert (yaml_modifier.start_epoch == serialized_modifier.start_epoch == obj_modifier.start_epoch)
    assert (sorted(yaml_modifier.submodules) == sorted(serialized_modifier.submodules) == sorted(obj_modifier.submodules))
    assert (yaml_modifier.model_fuse_fn_name == serialized_modifier.model_fuse_fn_name == obj_modifier.model_fuse_fn_name)
    assert (yaml_modifier.disable_quantization_observer_epoch == serialized_modifier.disable_quantization_observer_epoch == obj_modifier.disable_quantization_observer_epoch)
    assert (yaml_modifier.freeze_bn_stats_epoch == serialized_modifier.freeze_bn_stats_epoch == obj_modifier.freeze_bn_stats_epoch)
@pytest.mark.skipif(os.getenv('NM_ML_SKIP_PYTORCH_TESTS', False), reason='Skipping pytorch tests')
@pytest.mark.skipif(os.getenv('NM_ML_SKIP_PYTORCH_QUANT_TESTS', False), reason='Skipping pytorch torch quantization tests')
@pytest.mark.skipif((torch_quantization is None), reason='torch quantization not available')
def test_quantization_modifier_yaml():
    start_epoch = 0.0
    submodules = ['block.0', 'block.2']
    model_fuse_fn_name = 'fuse_module'
    disable_quantization_observer_epoch = 2.0
    freeze_bn_stats_epoch = 3.0
    yaml_str = '\n        !QuantizationModifier\n            start_epoch: {start_epoch}\n            submodules: {submodules}\n            model_fuse_fn_name: {model_fuse_fn_name}\n            disable_quantization_observer_epoch: {disable_quantization_observer_epoch}\n            freeze_bn_stats_epoch: {freeze_bn_stats_epoch}\n        '.format(start_epoch=start_epoch, submodules=submodules, model_fuse_fn_name=model_fuse_fn_name, disable_quantization_observer_epoch=disable_quantization_observer_epoch, freeze_bn_stats_epoch=freeze_bn_stats_epoch)
    yaml_modifier = QuantizationModifier.load_obj(yaml_str)
    serialized_modifier = QuantizationModifier.load_obj(str(yaml_modifier))
    obj_modifier = QuantizationModifier(start_epoch=start_epoch, submodules=submodules, model_fuse_fn_name=model_fuse_fn_name, disable_quantization_observer_epoch=disable_quantization_observer_epoch, freeze_bn_stats_epoch=freeze_bn_stats_epoch)
    assert isinstance(yaml_modifier, QuantizationModifier)
    assert (yaml_modifier.start_epoch == serialized_modifier.start_epoch == obj_modifier.start_epoch)
    assert (sorted(yaml_modifier.submodules) == sorted(serialized_modifier.submodules) == sorted(obj_modifier.submodules))
    assert (yaml_modifier.model_fuse_fn_name == serialized_modifier.model_fuse_fn_name == obj_modifier.model_fuse_fn_name)
    assert (yaml_modifier.disable_quantization_observer_epoch == serialized_modifier.disable_quantization_observer_epoch == obj_modifier.disable_quantization_observer_epoch)
    assert (yaml_modifier.freeze_bn_stats_epoch == serialized_modifier.freeze_bn_stats_epoch == obj_modifier.freeze_bn_stats_epoch)