def test_gm_pruning_yaml():
init_sparsity = 0.05
final_sparsity = 0.8
start_epoch = 5.0
end_epoch = 15.0
update_frequency = 1.0
params = ["re:.*weight"]
inter_func = "cubic"
mask_type = "filter"
global_sparsity = False
yaml_str = f"""
!GMPruningModifier
init_sparsity: {init_sparsity}
final_sparsity: {final_sparsity}
start_epoch: {start_epoch}
end_epoch: {end_epoch}
update_frequency: {update_frequency}
params: {params}
inter_func: {inter_func}
mask_type: {mask_type}
global_sparsity: {global_sparsity}
"""
yaml_modifier = GMPruningModifier.load_obj(
yaml_str) # type: GMPruningModifier
serialized_modifier = GMPruningModifier.load_obj(
str(yaml_modifier)) # type: GMPruningModifier
obj_modifier = GMPruningModifier(
init_sparsity=init_sparsity,
final_sparsity=final_sparsity,
start_epoch=start_epoch,
end_epoch=end_epoch,
update_frequency=update_frequency,
params=params,
inter_func=inter_func,
mask_type=mask_type,
global_sparsity=global_sparsity,
)
assert isinstance(yaml_modifier, GMPruningModifier)
assert (yaml_modifier.init_sparsity == serialized_modifier.init_sparsity ==
obj_modifier.init_sparsity)
assert (yaml_modifier.final_sparsity == serialized_modifier.final_sparsity
== obj_modifier.final_sparsity)
assert (yaml_modifier.start_epoch == serialized_modifier.start_epoch ==
obj_modifier.start_epoch)
assert (yaml_modifier.end_epoch == serialized_modifier.end_epoch ==
obj_modifier.end_epoch)
assert (yaml_modifier.update_frequency ==
serialized_modifier.update_frequency ==
obj_modifier.update_frequency)
assert yaml_modifier.params == serialized_modifier.params == obj_modifier.params
assert (yaml_modifier.inter_func == serialized_modifier.inter_func ==
obj_modifier.inter_func)
assert (str(yaml_modifier.mask_type) == str(serialized_modifier.mask_type)
== str(obj_modifier.mask_type))
assert (str(yaml_modifier.global_sparsity) == str(
serialized_modifier.global_sparsity) == str(
obj_modifier.global_sparsity))
def test_magnitude_pruning_yaml():
init_sparsity = 0.05
final_sparsity = 0.8
start_epoch = 5.0
end_epoch = 15.0
update_frequency = 1.0
params = "__ALL_PRUNABLE__"
inter_func = "cubic"
mask_type = "filter"
yaml_str = f"""
!MagnitudePruningModifier
init_sparsity: {init_sparsity}
final_sparsity: {final_sparsity}
start_epoch: {start_epoch}
end_epoch: {end_epoch}
update_frequency: {update_frequency}
params: {params}
inter_func: {inter_func}
mask_type: {mask_type}
"""
yaml_modifier = MagnitudePruningModifier.load_obj(
yaml_str) # type: MagnitudePruningModifier
serialized_modifier = GMPruningModifier.load_obj(
str(yaml_modifier)) # type: MagnitudePruningModifier
obj_modifier = GMPruningModifier(
init_sparsity=init_sparsity,
final_sparsity=final_sparsity,
start_epoch=start_epoch,
end_epoch=end_epoch,
update_frequency=update_frequency,
params=params,
inter_func=inter_func,
mask_type=mask_type,
)
assert isinstance(yaml_modifier, MagnitudePruningModifier)
assert (yaml_modifier.init_sparsity == serialized_modifier.init_sparsity ==
obj_modifier.init_sparsity)
assert (yaml_modifier.final_sparsity == serialized_modifier.final_sparsity
== obj_modifier.final_sparsity)
assert (yaml_modifier.start_epoch == serialized_modifier.start_epoch ==
obj_modifier.start_epoch)
assert (yaml_modifier.end_epoch == serialized_modifier.end_epoch ==
obj_modifier.end_epoch)
assert (yaml_modifier.update_frequency ==
serialized_modifier.update_frequency ==
obj_modifier.update_frequency)
assert yaml_modifier.params == serialized_modifier.params == obj_modifier.params
assert (yaml_modifier.inter_func == serialized_modifier.inter_func ==
obj_modifier.inter_func)
assert (str(yaml_modifier.mask_type) == str(serialized_modifier.mask_type)
== str(obj_modifier.mask_type))
def test_gm_pruning_yaml(params, final_sparsity):
init_sparsity = 0.05
start_epoch = 5.0
end_epoch = 15.0
update_frequency = 1.0
inter_func = "cubic"
mask_type = "filter"
global_sparsity = False
yaml_str = f"""
!LegacyGMPruningModifier
init_sparsity: {init_sparsity}
final_sparsity: {final_sparsity}
start_epoch: {start_epoch}
end_epoch: {end_epoch}
update_frequency: {update_frequency}
params: {params}
inter_func: {inter_func}
mask_type: {mask_type}
global_sparsity: {global_sparsity}
"""
yaml_modifier = GMPruningModifier.load_obj(yaml_str) # type: GMPruningModifier
serialized_modifier = GMPruningModifier.load_obj(
str(yaml_modifier)
) # type: GMPruningModifier
obj_modifier = GMPruningModifier(
init_sparsity=init_sparsity,
final_sparsity=final_sparsity,
start_epoch=start_epoch,
end_epoch=end_epoch,
update_frequency=update_frequency,
params=params,
inter_func=inter_func,
mask_type=mask_type,
global_sparsity=global_sparsity,
)
assert isinstance(yaml_modifier, GMPruningModifier)
_test_pruning_modifier_serialization_vals(
yaml_modifier, serialized_modifier, obj_modifier
)
assert (
str(yaml_modifier.global_sparsity)
== str(serialized_modifier.global_sparsity)
== str(obj_modifier.global_sparsity)
)
assert torch.all(param == 0.0)
@flaky(max_runs=3, min_passes=2)
@pytest.mark.skipif(
os.getenv("NM_ML_SKIP_PYTORCH_TESTS", False),
reason="Skipping pytorch tests",
)
@pytest.mark.parametrize(
"modifier_lambda",
[
lambda: GMPruningModifier(
init_sparsity=0.05,
final_sparsity=0.95,
start_epoch=0.0,
end_epoch=15.0,
update_frequency=1.0,
params=["re:.*weight"],
inter_func="linear",
),
lambda: GMPruningModifier(
params=["re:seq.block1.*weight"],
init_sparsity=0.05,
final_sparsity=0.95,
start_epoch=10.0,
end_epoch=25.0,
update_frequency=1.0,
inter_func="cubic",
global_sparsity=True,
),
lambda: GMPruningModifier(
obj_modifier.end_epoch)
assert yaml_modifier.params == serialized_modifier.params == obj_modifier.params
@pytest.mark.skipif(
os.getenv("NM_ML_SKIP_PYTORCH_TESTS", False),
reason="Skipping pytorch tests",
)
@pytest.mark.parametrize(
"modifier_lambda",
[
lambda: GMPruningModifier(
init_sparsity=0.05,
final_sparsity=0.95,
start_epoch=0.0,
end_epoch=15.0,
update_frequency=1.0,
params=["re:.*weight"],
inter_func="linear",
),
lambda: GMPruningModifier(
params=["re:seq.block1.*weight"],
init_sparsity=0.05,
final_sparsity=0.95,
start_epoch=10.0,
end_epoch=25.0,
update_frequency=1.0,
inter_func="cubic",
),
lambda: GMPruningModifier(
params=["seq.fc1.weight", "seq.fc2.weight"],
def create_config(project: Project, optim: ProjectOptimization,
framework: str) -> str:
"""
Creates a optimization config yaml for a given project and optimization
:param project: project to create with
:param optim: project optimizer to create with
:param framework: the framework to create the config for
"""
# add imports in function so they don't fail if they don't have env setup
# for frameworks other than the requested
if framework == "pytorch":
from sparseml.pytorch.optim import (
EpochRangeModifier,
GMPruningModifier,
LearningRateModifier,
ScheduledModifierManager,
SetLearningRateModifier,
TrainableParamsModifier,
)
elif framework == "tensorflow":
from sparseml.tensorflow_v1.optim import (
EpochRangeModifier,
GMPruningModifier,
LearningRateModifier,
ScheduledModifierManager,
SetLearningRateModifier,
TrainableParamsModifier,
)
else:
_LOGGER.error("Unsupported framework {} provided".format(framework))
raise ValidationError(
"Unsupported framework {} provided".format(framework))
mods = [
EpochRangeModifier(
start_epoch=optim.start_epoch
if optim.start_epoch is not None else -1,
end_epoch=optim.end_epoch if optim.end_epoch is not None else -1,
)
]
node_weight_name_lookup = {
node["id"]: node["weight_name"]
for node in project.model.analysis["nodes"] if node["prunable"]
}
for mod in optim.pruning_modifiers:
sparsity_to_params = {}
for node in mod.nodes:
# node is coming from DB, so already had prunable checks
# add assert here to fail early for non prunable nodes
assert node["node_id"] in node_weight_name_lookup
sparsity = node["sparsity"]
node_id = node["node_id"]
weight_name = node_weight_name_lookup[node_id]
if sparsity is None:
continue
if sparsity not in sparsity_to_params:
sparsity_to_params[sparsity] = []
sparsity_to_params[sparsity].append(weight_name)
for sparsity, params in sparsity_to_params.items():
gm_pruning = GMPruningModifier(
init_sparsity=0.05,
final_sparsity=sparsity,
start_epoch=mod.start_epoch
if mod.start_epoch is not None else -1,
end_epoch=mod.end_epoch if mod.end_epoch is not None else -1,
update_frequency=mod.update_frequency
if mod.update_frequency else -1,
params=params,
)
if mod.mask_type:
gm_pruning.mask_type = mod.mask_type
mods.append(gm_pruning)
for lr_schedule_modifier in optim.lr_schedule_modifiers:
for mod in lr_schedule_modifier.lr_mods:
mod = ProjectOptimizationModifierLRSchema().dump(mod)
start_epoch = mod["start_epoch"] if mod[
"start_epoch"] is not None else -1
end_epoch = mod["end_epoch"] if mod["end_epoch"] is not None else -1
if mod["clazz"] == "set":
mods.append(
SetLearningRateModifier(
learning_rate=mod["init_lr"],
start_epoch=start_epoch,
))
else:
lr_class_mapping = {
"step": "StepLR",
"multi_step": "MultiStepLR",
"exponential": "ExponentialLR",
}
assert mod["clazz"] in lr_class_mapping
mods.append(
LearningRateModifier(
lr_class=lr_class_mapping[mod["clazz"]],
lr_kwargs=mod["args"],
init_lr=mod["init_lr"],
start_epoch=start_epoch,
end_epoch=end_epoch,
))
for trainable_modifier in optim.trainable_modifiers:
mod = ProjectOptimizationModifierTrainableSchema().dump(
trainable_modifier)
start_epoch = mod["start_epoch"] if mod[
"start_epoch"] is not None else -1
end_epoch = mod["end_epoch"] if mod["end_epoch"] is not None else -1
if "nodes" not in mod:
continue
trainable_nodes = []
untrainable_nodes = []
for node in mod["nodes"]:
assert node["node_id"] in node_weight_name_lookup
weight_name = node_weight_name_lookup[node["node_id"]]
if node["trainable"]:
trainable_nodes.append(weight_name)
else:
untrainable_nodes.append(weight_name)
if len(trainable_nodes) > 0:
mods.append(
TrainableParamsModifier(trainable_nodes,
True,
start_epoch=start_epoch,
end_epoch=end_epoch))
if len(untrainable_nodes) > 0:
mods.append(
TrainableParamsModifier(
untrainable_nodes,
False,
start_epoch=start_epoch,
end_epoch=end_epoch,
))
# TODO: add quantization support when ready
return str(ScheduledModifierManager(mods))