Skip to content

Commit

Permalink
Update benchmark script for yolo
Browse files Browse the repository at this point in the history
osw282 committed Jan 29, 2025

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
1 parent 4ab1e62 commit 611da7a
Showing 10 changed files with 151,479 additions and 150,206 deletions.
61 changes: 43 additions & 18 deletions jetson/power_logging/model/benchmark_pruned_model.py
Original file line number Diff line number Diff line change
@@ -126,6 +126,30 @@ def layer_time_hook(layer_time_dict, layer_name, start_event, end_event, module,
layer_time_dict[layer_name]["start_time"] = start_event.get_time_stamp()


def get_layers_for_pruning(model: torch.nn.Module) -> tuple[(torch.nn.Module, str)]:
"""
Recursively get all layers in a pytorch model.
Args:
model: the pytorch model to look for layers.
name_prefix: Use to identify the parents layer. Defaults to "".
Returns:
a list of tuple containing the layer name and the layer.
"""
children = list(model.named_children())

if len(children) == 0:
result = [(model, "weight")]
else:
result = []
for _, child in children:
layers = get_layers_for_pruning(child)
result.extend(layers)

return tuple(result)


def benchmark(args: argparse.Namespace) -> None:
"""Benchmark latency and throughput across all backends.
@@ -138,24 +162,25 @@ def benchmark(args: argparse.Namespace) -> None:

try:
input_data = torch.randn(args.input_shape, device=DEVICE)
model = load_model(args.model, args.model_repo)
model.eval().to(DEVICE)

## PRUNING CODE FOR LENET ONLY
parameters_to_prune = (
(model.feat.conv1, 'weight'),
(model.feat.conv2, 'weight'),
(model.classifer.fc1, 'weight'),
(model.classifer.fc2, 'weight'),
(model.classifer.fc3, 'weight'),
)

prune.global_unstructured(
parameters_to_prune,
pruning_method=prune.L1Unstructured,
amount=0.3,
)

# model = load_model(args.model, args.model_repo)
# model.eval().to(DEVICE)

# ## PRUNING CODE FOR LENET ONLY
# parameters_to_prune = (
# (model.feat.conv1, 'weight'),
# (model.feat.conv2, 'weight'),
# (model.classifer.fc1, 'weight'),
# (model.classifer.fc2, 'weight'),
# (model.classifer.fc3, 'weight'),
# )

# prune.global_unstructured(
# parameters_to_prune,
# pruning_method=prune.L1Unstructured,
# amount=0.5,
# )
# Thus should load the pruned yolo model
model = torch.load("pruned_30_yolov5su.pt")
model.eval().to(DEVICE)

dtype = torch.float32
190 changes: 133 additions & 57 deletions jetson/power_logging/model_pruning_experiment.py
Original file line number Diff line number Diff line change
@@ -6,67 +6,143 @@
from model.model_utils import load_model


model = load_model("lenet", "pytorch/vision:v0.16.0")

parameters_to_prune = (
(model.feat.conv1, 'weight'),
(model.feat.conv2, 'weight'),
(model.classifer.fc1, 'weight'),
(model.classifer.fc2, 'weight'),
(model.classifer.fc3, 'weight'),
)
def get_layers_for_pruning(model: torch.nn.Module) -> tuple[(torch.nn.Module, str)]:
"""
Recursively get all layers in a pytorch model.
Args:
model: the pytorch model to look for layers.
name_prefix: Use to identify the parents layer. Defaults to "".
Returns:
a list of tuple containing the layer name and the layer.
"""
children = list(model.named_children())

if len(children) == 0:
result = [(model, "weight")]
else:
result = []
for _, child in children:
layers = get_layers_for_pruning(child)
result.extend(layers)

return tuple(result)

# Lenet experiment
# model = load_model("lenet", "pytorch/vision:v0.16.0")

# parameters_to_prune = (
# (model.feat.conv1, 'weight'),
# (model.feat.conv2, 'weight'),
# (model.classifer.fc1, 'weight'),
# (model.classifer.fc2, 'weight'),
# (model.classifer.fc3, 'weight'),
# )

# prune.global_unstructured(
# parameters_to_prune,
# pruning_method=prune.L1Unstructured,
# amount=0.3,
# )

# print(
# "Sparsity in conv1.weight: {:.2f}%".format(
# 100. * float(torch.sum(model.feat.conv1.weight == 0))
# / float(model.feat.conv1.weight.nelement())
# )
# )
# print(
# "Sparsity in conv2.weight: {:.2f}%".format(
# 100. * float(torch.sum(model.feat.conv2.weight == 0))
# / float(model.feat.conv2.weight.nelement())
# )
# )
# print(
# "Sparsity in fc1.weight: {:.2f}%".format(
# 100. * float(torch.sum(model.classifer.fc1.weight == 0))
# / float(model.classifer.fc1.weight.nelement())
# )
# )
# print(
# "Sparsity in fc2.weight: {:.2f}%".format(
# 100. * float(torch.sum(model.classifer.fc2.weight == 0))
# / float(model.classifer.fc2.weight.nelement())
# )
# )
# print(
# "Sparsity in fc3.weight: {:.2f}%".format(
# 100. * float(torch.sum(model.classifer.fc3.weight == 0))
# / float(model.classifer.fc3.weight.nelement())
# )
# )
# print(
# "Global sparsity: {:.2f}%".format(
# 100. * float(
# torch.sum(model.feat.conv1.weight == 0)
# + torch.sum(model.feat.conv2.weight == 0)
# + torch.sum(model.classifer.fc1.weight == 0)
# + torch.sum(model.classifer.fc2.weight == 0)
# + torch.sum(model.classifer.fc3.weight == 0)
# )
# / float(
# model.feat.conv1.weight.nelement()
# + model.feat.conv2.weight.nelement()
# + model.classifer.fc1.weight.nelement()
# + model.classifer.fc2.weight.nelement()
# + model.classifer.fc3.weight.nelement()
# )
# )
# )



# # VGG19
# model = load_model("vgg19", "pytorch/vision:v0.16.0")

# parameters_to_prune = get_layers_for_pruning(model)

# parameters_to_prune = [(layer, "weight") for layer, _ in parameters_to_prune if isinstance(layer, torch.nn.modules.conv.Conv2d)\
# or isinstance(layer, torch.nn.modules.pooling.MaxPool2d)\
# or isinstance(layer, torch.nn.modules.linear.Linear)]

# prune.global_unstructured(
# parameters_to_prune,
# pruning_method=prune.L1Unstructured,
# amount=0.5,
# )

#
# YOLOV5 MODEL DETECTION
from ultralytics import YOLO
# model = torch.hub.load('ultralytics/yolov5', 'yolov5su', pretrained=True)

model = YOLO("yolov5su.pt")
parameters_to_prune = get_layers_for_pruning(model)
parameters_to_prune = [(layer, "weight") for layer, _ in parameters_to_prune if isinstance(layer, torch.nn.modules.conv.Conv2d) or isinstance(layer, torch.nn.modules.linear.Linear)]
# print(parameters_to_prune)

prune.global_unstructured(
parameters_to_prune,
pruning_method=prune.L1Unstructured,
amount=0.3,
)
model.save("pruned_30_yolov5su.pt")

print(
"Sparsity in conv1.weight: {:.2f}%".format(
100. * float(torch.sum(model.feat.conv1.weight == 0))
/ float(model.feat.conv1.weight.nelement())
)
)
print(
"Sparsity in conv2.weight: {:.2f}%".format(
100. * float(torch.sum(model.feat.conv2.weight == 0))
/ float(model.feat.conv2.weight.nelement())
)
)
print(
"Sparsity in fc1.weight: {:.2f}%".format(
100. * float(torch.sum(model.classifer.fc1.weight == 0))
/ float(model.classifer.fc1.weight.nelement())
)
)
print(
"Sparsity in fc2.weight: {:.2f}%".format(
100. * float(torch.sum(model.classifer.fc2.weight == 0))
/ float(model.classifer.fc2.weight.nelement())
)
)
print(
"Sparsity in fc3.weight: {:.2f}%".format(
100. * float(torch.sum(model.classifer.fc3.weight == 0))
/ float(model.classifer.fc3.weight.nelement())
)
)
print(
"Global sparsity: {:.2f}%".format(
100. * float(
torch.sum(model.feat.conv1.weight == 0)
+ torch.sum(model.feat.conv2.weight == 0)
+ torch.sum(model.classifer.fc1.weight == 0)
+ torch.sum(model.classifer.fc2.weight == 0)
+ torch.sum(model.classifer.fc3.weight == 0)
)
/ float(
model.feat.conv1.weight.nelement()
+ model.feat.conv2.weight.nelement()
+ model.classifer.fc1.weight.nelement()
+ model.classifer.fc2.weight.nelement()
+ model.classifer.fc3.weight.nelement()
)
)
)
# print(list(parameters_to_prune[0][0].named_buffers()))
# model = torch.load("pruned_yolov5su.pt") # local model
model = YOLO("pruned_30_yolov5su.pt")
# Run evaluation on COCO dataset
results = model.val(data="coco.yaml")


# model = torch.load("yolov5su.pt") # local model, yolo weight
# parameters_to_prune = get_layers_for_pruning(model['model'])
# parameters_to_prune = [(layer, "weight") for layer, _ in parameters_to_prune if isinstance(layer, torch.nn.modules.conv.Conv2d) or isinstance(layer, torch.nn.modules.linear.Linear)]

# prune.global_unstructured(
# parameters_to_prune,
# pruning_method=prune.L1Unstructured,
# amount=0.3,
# )
# model.save("pruned_yolov5su.pt")
2 changes: 2 additions & 0 deletions jetson/power_logging/pyproject.toml
Original file line number Diff line number Diff line change
@@ -7,9 +7,11 @@ requires-python = ">=3.11"
dependencies = [
"boto3<1.36",
"dvc-s3>=3.2.0",
"opencv-python>=4.11.0.86",
"pandas>=2.2.3",
"pillow>=11.1.0",
"torch==2.2.2",
"ultralytics>=8.3.69",
]

[dependency-groups]
Loading

0 comments on commit 611da7a

Please sign in to comment.