def get_basic_sparsity_config(model_size=4,
input_sample_size=None,
sparsity_init=0.02,
sparsity_target=0.5,
sparsity_target_epoch=2,
sparsity_freeze_epoch=3,
scheduler='polinomial'):
if input_sample_size is None:
input_sample_size = [1, 1, 4, 4]
config = NNCFConfig()
config.update({
"model": "basic_sparse_conv",
"model_size": model_size,
"input_info": {
"sample_size": input_sample_size,
},
"compression": {
"algorithm": "rb_sparsity",
"sparsity_init": sparsity_init,
"params": {
"schedule": scheduler,
"sparsity_target": sparsity_target,
"sparsity_target_epoch": sparsity_target_epoch,
"sparsity_freeze_epoch": sparsity_freeze_epoch
},
}
})
return config
def test_model_device_before_create_compressed_model(device_placing,
inference_type):
if not torch.cuda.is_available() and not inference_type == 'cpu':
pytest.skip("Skipping CUDA test cases for CPU only setups")
input_size = [1, 1, 8, 8]
config = NNCFConfig()
config = get_kd_config(config)
config.update({
"input_info": {
"sample_size": input_size,
},
})
if inference_type == 'DDP':
ngpus_per_node = torch.cuda.device_count()
config.world_size = ngpus_per_node
torch.multiprocessing.spawn(run_training_for_device_testing,
nprocs=ngpus_per_node,
args=(config, inference_type,
ngpus_per_node, device_placing),
join=True)
else:
run_training_for_device_testing(None,
config,
inference_type,
None,
device_placing=device_placing)
def worker(rank: int, world_size: int) -> None:
torch.distributed.init_process_group(backend="nccl",
init_method='tcp://127.0.0.1:8999',
world_size=world_size,
rank=rank)
model = TestModelWithChangedTrain(freezing_stages=1)
model.cuda()
model.to(rank)
nncf_config = NNCFConfig()
nncf_config.update({
"input_info": {
"sample_size": [1, 1, 30, 30]
},
"compression": {
"algorithm": "quantization",
"initializer": {
"range": {
"num_init_samples": 10
},
"batchnorm_adaptation": {
"num_bn_adaptation_samples": 10
}
}
}
})
dataloader = create_random_mock_dataloader(nncf_config, num_samples=10)
register_default_init_args(nncf_config, dataloader)
_, compressed_model = create_compressed_model(model, nncf_config)
# At this part the additional processes may be freezing
_ = torch.nn.parallel.DistributedDataParallel(compressed_model,
device_ids=[rank])
def get_basic_sparsity_config(model_size=4,
input_sample_size=None,
sparsity_init=0.02,
sparsity_target=0.5,
sparsity_target_epoch=2,
sparsity_freeze_epoch=3):
if input_sample_size is None:
input_sample_size = [1, 1, 4, 4]
config = NNCFConfig()
config.update({
'model': 'basic_sparse_conv',
'model_size': model_size,
'input_info': {
'sample_size': input_sample_size,
},
'compression': {
'algorithm': 'rb_sparsity',
'sparsity_init': sparsity_init,
'params': {
'schedule': 'polynomial',
'sparsity_target': sparsity_target,
'sparsity_target_epoch': sparsity_target_epoch,
'sparsity_freeze_epoch': sparsity_freeze_epoch
},
}
})
return config
def get_config_for_logarithm_scale(logarithm_scale: bool,
quantization_type: str) -> NNCFConfig:
nncf_config = NNCFConfig()
nncf_config.update({
"input_info": {
"sample_size": SAMPLE_SIZE
},
"target_device": 'TRIAL',
"compression": {
"algorithm": "quantization",
"initializer": {
"range": {
"num_init_samples": 4,
"type": "percentile",
"params": {
"min_percentile": 0.001,
"max_percentile": 99.999
}
}
},
"activations": {
"mode": quantization_type,
"logarithm_scale": logarithm_scale
},
"weights": {
"mode": quantization_type,
"signed": True,
"logarithm_scale": logarithm_scale
}
}
})
class RandDatasetMock:
def __getitem__(self, index):
return torch.rand(*SAMPLE_SIZE)
def __len__(self):
return 4
data_loader = torch.utils.data.DataLoader(RandDatasetMock(),
batch_size=1,
shuffle=False,
drop_last=True)
class SquadInitializingDataloader(
nncf.torch.initialization.PTInitializingDataLoader):
def get_inputs(self, batch):
return batch, {}
def get_target(self, batch):
return None
initializing_data_loader = SquadInitializingDataloader(data_loader)
init_range = nncf.config.structures.QuantizationRangeInitArgs(
initializing_data_loader)
nncf_config.register_extra_structs([init_range])
register_bn_adaptation_init_args(nncf_config)
return nncf_config
def get_config_for_export_mode(should_be_onnx_standard: bool) -> NNCFConfig:
nncf_config = NNCFConfig()
nncf_config.update({
"input_info": {
"sample_size": [1, 1, 4, 4]
},
"compression": {
"algorithm": "quantization",
"export_to_onnx_standard_ops": should_be_onnx_standard
}
})
return nncf_config
def _get_mock_config(algo_name: Union[List[str], str]) -> NNCFConfig:
config = NNCFConfig()
config["input_info"] = {"sample_size": [1, 1]}
if isinstance(algo_name, list):
lst = []
for alg_n in algo_name:
lst.append({"algorithm": alg_n})
config["compression"] = lst
else:
assert isinstance(algo_name, str)
config["compression"] = {"algorithm": algo_name}
return config
def create_test_quantization_env(model_creator=BasicConvTestModel,
input_info_cfg=None) -> QuantizationEnv:
if input_info_cfg is None:
input_info_cfg = {"input_info": {"sample_size": [1, 1, 4, 4]}}
model = model_creator()
nncf_network = NNCFNetwork(model,
input_infos=create_input_infos(input_info_cfg))
hw_config_type = HWConfigType.VPU
hw_config_path = HWConfig.get_path_to_hw_config(hw_config_type)
hw_config = PTHWConfig.from_json(hw_config_path)
setup = PropagationBasedQuantizerSetupGenerator(
NNCFConfig(), nncf_network, hw_config=hw_config).generate_setup()
dummy_multi_setup = MultiConfigQuantizerSetup.from_single_config_setup(
setup)
for qp in dummy_multi_setup.quantization_points.values():
qconf_constraint_list = []
qconf = qp.possible_qconfigs[0]
bit_set = [8, 4, 2] if 'conv' in str(qp.insertion_point) else [8, 4]
for bits in bit_set:
adj_qconf = deepcopy(qconf)
adj_qconf.num_bits = bits
qconf_constraint_list.append(adj_qconf)
qp.possible_qconfigs = qconf_constraint_list
experimental_builder = ExperimentalQuantizationBuilder(
dummy_multi_setup, setup, {}, hw_config)
experimental_builder.apply_to(nncf_network)
# pylint:disable=line-too-long
experimental_ctrl = experimental_builder.build_controller(nncf_network)
data_loader = create_ones_mock_dataloader(input_info_cfg)
constraints = HardwareQuantizationConstraints()
for qid, qp_id_set in experimental_ctrl.module_id_to_qp_id_translation_dict.items(
):
first_qp_id_for_this_quantizer_module = next(iter(qp_id_set))
qconfigs = dummy_multi_setup.quantization_points[
first_qp_id_for_this_quantizer_module].possible_qconfigs
constraints.add(qid, qconfigs)
return QuantizationEnv(nncf_network,
experimental_ctrl,
constraints,
data_loader,
lambda *x: 0,
hw_config_type=HWConfigType.VPU,
params=QuantizationEnvParams(
compression_ratio=0.15,
eval_subset_ratio=1.0,
skip_constraint=False,
performant_bw=False,
finetune=False,
bits=[2, 4, 8],
dump_init_precision_data=False))
def get_basic_quantization_config(model_size=4):
config = NNCFConfig()
config.update(
Dict({
'model': 'basic_quant_conv',
'input_info': {
'sample_size': [1, model_size, model_size, 1],
},
'compression': {
'algorithm': 'quantization',
}
}))
return config
def get_basic_filter_pruning_config(input_sample_size=None):
if input_sample_size is None:
input_sample_size = [1, 4, 4, 1]
config = NNCFConfig({
"model": "basic_prune_conv",
"input_info": {
"sample_size": input_sample_size,
},
"compression": {
"algorithm": "filter_pruning",
"params": {}
}
})
return config
def get_basic_magnitude_sparsity_config(input_sample_size=None):
if input_sample_size is None:
input_sample_size = [1, 4, 4, 1]
config = NNCFConfig({
"model": "basic_sparse_conv",
"input_info": {
"sample_size": input_sample_size,
},
"compression": {
"algorithm": "magnitude_sparsity",
"params": {}
}
})
return config
def get_empty_config(input_sample_sizes=None) -> NNCFConfig:
if input_sample_sizes is None:
input_sample_sizes = [1, 4, 4, 1]
def _create_input_info():
if isinstance(input_sample_sizes, tuple):
return [{"sample_size": sizes} for sizes in input_sample_sizes]
return [{"sample_size": input_sample_sizes}]
config = NNCFConfig({
"model": "basic_sparse_conv",
"input_info": _create_input_info()
})
return config
def test_scheduler_can_do_epoch_step__with_rb_algo():
config = NNCFConfig()
config['input_info'] = [{"sample_size": [1, 4, 4, 1]}]
config['compression'] = {
'algorithm': 'rb_sparsity',
'sparsity_init': 0.2,
"params": {
'schedule': 'polynomial',
'power': 1,
'sparsity_target_epoch': 2,
'sparsity_target': 0.6,
'sparsity_freeze_epoch': 3
}
}
_, compression_ctrl = create_compressed_model_and_algo_for_test(
get_basic_conv_test_model(), config)
scheduler = compression_ctrl.scheduler
loss = compression_ctrl.loss
assert not loss.disabled
# pylint: disable=protected-access
for op, op_weights in loss._target_ops:
assert op.get_trainable_weight(op_weights)
scheduler.epoch_step()
assert pytest.approx(loss.target_sparsity_rate, abs=1e-3) == 0.2
assert pytest.approx(loss(), abs=1e-3) == 16
assert not loss.disabled
scheduler.epoch_step()
assert pytest.approx(loss.target_sparsity_rate, abs=1e-3) == 0.4
assert pytest.approx(loss(), abs=1e-3) == 64
assert not loss.disabled
scheduler.epoch_step()
assert pytest.approx(loss.target_sparsity_rate, abs=1e-3) == 0.6
assert pytest.approx(loss(), abs=1e-3) == 144
assert not loss.disabled
scheduler.epoch_step()
assert loss.disabled
assert pytest.approx(loss.target_sparsity_rate, abs=1e-3) == 0.6
assert loss() == 0
for op, op_weights in loss._target_ops:
assert not op.get_trainable_weight(op_weights)
def get_quantization_config_without_range_init(model_size=4):
config = NNCFConfig()
config.update({
"model": "basic_quant_conv",
"model_size": model_size,
"input_info": {
"sample_size": [1, 1, model_size, model_size],
},
"compression": {
"algorithm": "quantization",
"initializer": {
"range": {
"num_init_samples": 0
}
}
}
})
return config
def get_basic_pruning_config(model_size=8):
config = NNCFConfig()
config.update(Dict({
"model": "basic",
"input_info":
{
"sample_size": [1, model_size, model_size, 1],
},
"compression":
{
"algorithm": "filter_pruning",
"pruning_init": 0.5,
"params": {
"prune_first_conv": True,
}
}
}))
return config
def get_config_for_test(batch_size=10, num_bn_adaptation_samples=100):
config = NNCFConfig()
config.update(
Dict({
"compression": {
"algorithm": "quantization",
"initializer": {
"batchnorm_adaptation": {
"num_bn_adaptation_samples": num_bn_adaptation_samples,
}
}
}
}))
dataset = get_dataset_for_test()
config = register_default_init_args(config, dataset, batch_size)
return config
def get_binarization_config() -> NNCFConfig:
config = NNCFConfig()
config.update({
"model":
"resnet18",
"input_info": {
"sample_size": [1, *LeNet.INPUT_SIZE]
},
"compression": [{
"algorithm": "binarization",
"mode": "xnor",
"params": {
"activations_quant_start_epoch": 0,
"weights_quant_start_epoch": 0
}
}]
})
return config
def get_basic_quantization_config():
config = NNCFConfig()
config.update({
"model": "AlexNet",
"input_info": {
"sample_size": [1, 3, 32, 32],
},
"compression": {
"algorithm": "quantization",
"quantize_inputs": True,
"initializer": {
"range": {
"num_init_samples": 0
}
}
}
})
return config
def get_basic_quantization_config():
config = NNCFConfig()
config.update({
'model': 'AlexNet',
'input_info': {
'sample_size': [1, 3, 32, 32],
},
'compression': {
'algorithm': 'quantization',
'quantize_inputs': True,
'initializer': {
'range': {
'num_init_samples': 0
}
}
}
})
register_bn_adaptation_init_args(config)
return config
def create_test_quantization_env() -> QuantizationEnv:
model = BasicConvTestModel()
nncf_network = NNCFNetwork(model,
input_infos=[ModelInputInfo([1, 1, 4, 4])])
hw_config_type = HWConfigType.VPU
hw_config_path = HWConfig.get_path_to_hw_config(hw_config_type)
hw_config = HWConfig.from_json(hw_config_path)
setup = PropagationBasedQuantizerSetupGenerator(
NNCFConfig(), nncf_network, hw_config=hw_config).generate_setup()
experimental_builder = ExperimentalQuantizationBuilder(setup, {})
experimental_builder.apply_to(nncf_network)
# pylint:disable=line-too-long
experimental_ctrl = nncf_network.commit_compression_changes(
) # type: ExperimentalQuantizationController
data_loader = create_mock_dataloader({
"sample_size": [1, 1, 4, 4],
})
constraints = HardwareQuantizationConstraints()
for qid in experimental_ctrl.all_quantizations:
qconf_constraint_list = []
qconf = experimental_ctrl.all_quantizations[qid].get_current_config()
bit_set = [8, 4, 2] if 'conv' in str(qid) else [8, 4]
for bits in bit_set:
adj_qconf = deepcopy(qconf)
adj_qconf.bits = bits
qconf_constraint_list.append(adj_qconf)
constraints.add(qid, qconf_constraint_list)
return QuantizationEnv(nncf_network,
experimental_ctrl,
constraints,
data_loader,
lambda *x: 0,
hw_config_type=HWConfigType.VPU,
params=QuantizationEnvParams(
compression_ratio=0.15,
eval_subset_ratio=1.0,
skip_constraint=False,
finetune=False,
bits=[2, 4, 8],
dump_init_precision_data=False))
def wrap_nncf_model(model,
cfg,
data_loader_for_init=None,
get_fake_input_func=None):
"""
The function wraps mmdet model by NNCF
Note that the parameter `get_fake_input_func` should be the function `get_fake_input`
-- cannot import this function here explicitly
"""
check_nncf_is_enabled()
pathlib.Path(cfg.work_dir).mkdir(parents=True, exist_ok=True)
nncf_config = NNCFConfig(cfg.nncf_config)
logger = get_root_logger(cfg.log_level)
if data_loader_for_init:
wrapped_loader = MMInitializeDataLoader(data_loader_for_init)
nncf_config = register_default_init_args(nncf_config, None,
wrapped_loader)
if cfg.get('resume_from'):
checkpoint_path = cfg.get('resume_from')
assert is_checkpoint_nncf(checkpoint_path), (
'It is possible to resume training with NNCF compression from NNCF checkpoints only. '
'Use "load_from" with non-compressed model for further compression by NNCF.'
)
elif cfg.get('load_from'):
checkpoint_path = cfg.get('load_from')
if not is_checkpoint_nncf(checkpoint_path):
checkpoint_path = None
logger.info('Received non-NNCF checkpoint to start training '
'-- initialization of NNCF fields will be done')
else:
checkpoint_path = None
if not data_loader_for_init and not checkpoint_path:
raise RuntimeError('Either data_loader_for_init or NNCF pre-trained '
'model checkpoint should be set')
if checkpoint_path:
logger.info(f'Loading NNCF checkpoint from {checkpoint_path}')
logger.info(
'Please, note that this first loading is made before addition of '
'NNCF FakeQuantize nodes to the model, so there may be some '
'warnings on unexpected keys')
resuming_state_dict = load_checkpoint(model, checkpoint_path)
logger.info(f'Loaded NNCF checkpoint from {checkpoint_path}')
else:
resuming_state_dict = None
if "nncf_compress_postprocessing" in cfg:
# NB: This parameter is used to choose if we should try to make NNCF compression
# for a whole model graph including postprocessing (`nncf_compress_postprocessing=True`),
# or make NNCF compression of the part of the model without postprocessing
# (`nncf_compress_postprocessing=False`).
# Our primary goal is to make NNCF compression of such big part of the model as
# possible, so `nncf_compress_postprocessing=True` is our primary choice, whereas
# `nncf_compress_postprocessing=False` is our fallback decision.
# When we manage to enable NNCF compression for sufficiently many models,
# we should keep one choice only.
nncf_compress_postprocessing = cfg.get('nncf_compress_postprocessing')
logger.debug('set should_compress_postprocessing='
f'{nncf_compress_postprocessing}')
else:
nncf_compress_postprocessing = True
def _get_fake_data_for_forward(cfg, nncf_config, get_fake_input_func):
input_size = nncf_config.get("input_info").get('sample_size')
assert get_fake_input_func is not None
assert len(input_size) == 4 and input_size[0] == 1
H, W, C = input_size[2], input_size[3], input_size[1]
device = next(model.parameters()).device
return get_fake_input_func(cfg,
orig_img_shape=tuple([H, W, C]),
device=device)
def dummy_forward(model):
fake_data = _get_fake_data_for_forward(cfg, nncf_config,
get_fake_input_func)
img, img_metas = fake_data["img"], fake_data["img_metas"]
img = nncf_model_input(img)
if nncf_compress_postprocessing:
ctx = model.forward_export_context(img_metas)
logger.debug(
f"NNCF will compress a postprocessing part of the model")
else:
ctx = model.forward_dummy_context(img_metas)
logger.debug(
f"NNCF will NOT compress a postprocessing part of the model")
with ctx:
model(img)
model.dummy_forward_fn = dummy_forward
compression_ctrl, model = create_compressed_model(
model,
nncf_config,
dummy_forward_fn=dummy_forward,
resuming_state_dict=resuming_state_dict)
model = change_export_func_first_conv(model)
return compression_ctrl, model
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
from nncf import NNCFConfig
from tests.tensorflow.helpers import create_compressed_model_and_algo_for_test
from tests.tensorflow.helpers import get_basic_two_conv_test_model
from tests.tensorflow.helpers import TFTensorListComparator
EPS = 1e-9
INPUT_SIZE = [4, 4, 1]
NO_COMPRESSION_NNCF_CONFIG = NNCFConfig({
"model": "basic_binarization_config",
"input_info": {
"sample_size": [1] + INPUT_SIZE
}
})
def test_no_compression_algo_not_change_model_params():
orig_model = get_basic_two_conv_test_model()
model, _algo = create_compressed_model_and_algo_for_test(
orig_model, NO_COMPRESSION_NNCF_CONFIG)
orig_model_weights = orig_model.get_weights()
model_weights = model.get_weights()
TFTensorListComparator.check_equal(orig_model_weights, model_weights)
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"--epochs", type=int, default=50, help="number of epochs to train (default: 50)"
)
parser.add_argument(
"--lr", type=float, default=0.05, help="learning rate (default: 0.05)"
)
parser.add_argument(
"--enable_nncf_compression",
action="store_true",
default=False,
help="nncf compression flag (default: False)",
)
parser.add_argument("--seed", type=int, default=1, help="random seed (default: 1)")
parser.add_argument(
"--ckpt_filename",
type=str,
default="resnet18_cifar10.pth",
help="file name for model checkpoint (default: resnet18_cifar10.pth)",
)
parser.add_argument(
"--starting_checkpoint",
type=str,
default=None,
help="checkpoint file name to start training from (default: None)",
)
args = parser.parse_args()
print(args)
torch.manual_seed(args.seed)
input_size, num_classes, train_dataset, test_dataset = get_CIFAR10()
kwargs = {"num_workers": 8, "pin_memory": True}
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=512, shuffle=True, **kwargs
)
test_loader = torch.utils.data.DataLoader(
test_dataset, batch_size=5000, shuffle=False, **kwargs
)
model = Model()
model = model.cuda()
if args.starting_checkpoint is not None:
model.load_state_dict(torch.load(args.starting_checkpoint))
compression_ctrl = None
if args.enable_nncf_compression:
nncf_config_dict = {
"compression": {
"algorithm": "quantization",
"initializer": {"range": {"num_init_steps": 5}},
}
}
nncf_config = NNCFConfig(nncf_config_dict)
nncf_config = register_default_init_args(nncf_config, None, train_loader)
compression_ctrl, model = create_compressed_model(model, nncf_config)
if args.enable_nncf_compression:
milestones = [5, 10]
else:
milestones = [25, 40]
optimizer = torch.optim.SGD(
model.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4
)
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=milestones, gamma=0.1
)
for epoch in range(1, args.epochs + 1):
train(model, train_loader, optimizer, epoch, compression_ctrl)
test(model, test_loader)
scheduler.step()
if compression_ctrl is not None:
compression_ctrl.scheduler.epoch_step()
torch.save(model.state_dict(), args.ckpt_filename)
def wrap_nncf_model(model,
cfg,
checkpoint_dict=None,
datamanager_for_init=None):
# Note that we require to import it here to avoid cyclic imports when import get_no_nncf_trace_context_manager
# from mobilenetv3
from torchreid.data.transforms import build_inference_transform
from nncf import NNCFConfig
from nncf.torch import create_compressed_model, load_state
from nncf.torch.initialization import register_default_init_args
from nncf.torch.dynamic_graph.io_handling import nncf_model_input
from nncf.torch.dynamic_graph.trace_tensor import TracedTensor
from nncf.torch.initialization import PTInitializingDataLoader
if checkpoint_dict is None:
checkpoint_path = cfg.model.load_weights
resuming_checkpoint = safe_load_checkpoint(
checkpoint_path, map_location=torch.device('cpu'))
else:
checkpoint_path = 'pretrained_dict'
resuming_checkpoint = checkpoint_dict
if datamanager_for_init is None and not is_nncf_state(resuming_checkpoint):
raise RuntimeError('Either datamanager_for_init or NNCF pre-trained '
'model checkpoint should be set')
nncf_metainfo = None
if is_nncf_state(resuming_checkpoint):
nncf_metainfo = _get_nncf_metainfo_from_state(resuming_checkpoint)
nncf_config_data = nncf_metainfo['nncf_config']
datamanager_for_init = None
logger.info(f'Read NNCF metainfo with NNCF config from the checkpoint:'
f'nncf_metainfo=\n{pformat(nncf_metainfo)}')
else:
resuming_checkpoint = None
nncf_config_data = cfg.get('nncf_config')
if nncf_config_data is None:
logger.info('Cannot read nncf_config from config file')
else:
logger.info(f' nncf_config=\n{pformat(nncf_config_data)}')
h, w = cfg.data.height, cfg.data.width
if not nncf_config_data:
logger.info('Using the default NNCF int8 quantization config')
nncf_config_data = get_default_nncf_compression_config(h, w)
# do it even if nncf_config_data is loaded from a checkpoint -- for the rare case when
# the width and height of the model's input was changed in the config
# and then finetuning of NNCF model is run
nncf_config_data.setdefault('input_info', {})
nncf_config_data['input_info']['sample_size'] = [1, 3, h, w]
nncf_config = NNCFConfig(nncf_config_data)
logger.info(f'nncf_config =\n{pformat(nncf_config)}')
if not nncf_metainfo:
nncf_metainfo = create_nncf_metainfo(enable_quantization=True,
enable_pruning=False,
nncf_config=nncf_config_data)
else:
# update it just to be on the safe side
nncf_metainfo['nncf_config'] = nncf_config_data
class ReidInitializeDataLoader(PTInitializingDataLoader):
def get_inputs(self, dataloader_output):
# define own InitializingDataLoader class using approach like
# parse_data_for_train and parse_data_for_eval in the class Engine
# dataloader_output[0] should be image here
args = (dataloader_output[0], )
return args, {}
@torch.no_grad()
def model_eval_fn(model):
"""
Runs evaluation of the model on the validation set and
returns the target metric value.
Used to evaluate the original model before compression
if NNCF-based accuracy-aware training is used.
"""
from torchreid.metrics.classification import evaluate_classification
if test_loader is None:
raise RuntimeError(
'Cannot perform a model evaluation on the validation '
'dataset since the validation data loader was not passed '
'to wrap_nncf_model')
model_type = get_model_attr(model, 'type')
targets = list(test_loader.keys())
use_gpu = cur_device.type == 'cuda'
for dataset_name in targets:
domain = 'source' if dataset_name in datamanager_for_init.sources else 'target'
print(f'##### Evaluating {dataset_name} ({domain}) #####')
if model_type == 'classification':
cmc, _, _ = evaluate_classification(
test_loader[dataset_name]['query'], model, use_gpu=use_gpu)
accuracy = cmc[0]
elif model_type == 'multilabel':
mAP, _, _, _, _, _, _ = evaluate_multilabel_classification(
test_loader[dataset_name]['query'], model, use_gpu=use_gpu)
accuracy = mAP
else:
raise ValueError(
f'Cannot perform a model evaluation on the validation dataset'
f'since the model has unsupported model_type {model_type or "None"}'
)
return accuracy
cur_device = next(model.parameters()).device
logger.info(f'NNCF: cur_device = {cur_device}')
if resuming_checkpoint is None:
logger.info(
'No NNCF checkpoint is provided -- register initialize data loader'
)
train_loader = datamanager_for_init.train_loader
test_loader = datamanager_for_init.test_loader
wrapped_loader = ReidInitializeDataLoader(train_loader)
nncf_config = register_default_init_args(nncf_config,
wrapped_loader,
model_eval_fn=model_eval_fn,
device=cur_device)
model_state_dict = None
compression_state = None
else:
model_state_dict, compression_state = extract_model_and_compression_states(
resuming_checkpoint)
transform = build_inference_transform(
cfg.data.height,
cfg.data.width,
norm_mean=cfg.data.norm_mean,
norm_std=cfg.data.norm_std,
)
def dummy_forward(model):
prev_training_state = model.training
model.eval()
input_img = random_image(cfg.data.height, cfg.data.width)
input_blob = transform(input_img).unsqueeze(0)
assert len(input_blob.size()) == 4
input_blob = input_blob.to(device=cur_device)
input_blob = nncf_model_input(input_blob)
model(input_blob)
model.train(prev_training_state)
def wrap_inputs(args, kwargs):
assert len(args) == 1
if isinstance(args[0], TracedTensor):
logger.info('wrap_inputs: do not wrap input TracedTensor')
return args, {}
return (nncf_model_input(args[0]), ), kwargs
model.dummy_forward_fn = dummy_forward
if 'log_dir' in nncf_config:
os.makedirs(nncf_config['log_dir'], exist_ok=True)
logger.info(f'nncf_config["log_dir"] = {nncf_config["log_dir"]}')
compression_ctrl, model = create_compressed_model(
model,
nncf_config,
dummy_forward_fn=dummy_forward,
wrap_inputs_fn=wrap_inputs,
compression_state=compression_state)
if model_state_dict:
logger.info(f'Loading NNCF model from {checkpoint_path}')
load_state(model, model_state_dict, is_resume=True)
return compression_ctrl, model, nncf_metainfo
