From 504adbe552ec9aabf1e9693f1f9033f8ed152ab8 Mon Sep 17 00:00:00 2001
From: spencerwooo <spencer.woo@outlook.com>
Date: Mon, 15 Apr 2024 12:50:29 +0800
Subject: [PATCH] Add multiple layer extract example

---
 examples/extract_multiple_layers.py | 21 +++++++++++++++++++++
 requirements.txt                    |  2 +-
 2 files changed, 22 insertions(+), 1 deletion(-)
 create mode 100644 examples/extract_multiple_layers.py

diff --git a/examples/extract_multiple_layers.py b/examples/extract_multiple_layers.py
new file mode 100644
index 0000000..5882b97
--- /dev/null
+++ b/examples/extract_multiple_layers.py
@@ -0,0 +1,21 @@
+import torch
+from torchvision.models import resnet50
+
+from torch_featurelayer import FeatureLayers
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Load a pretrained ResNet-50 model
+model = resnet50(weights='DEFAULT').eval().to(device)
+
+# Hook onto layers ['layer1', 'layer2', 'layer3', 'layer4'] of the model
+layer_paths = ['layer1', 'layer2', 'layer3', 'layer4']
+hooked_model = FeatureLayers(model, layer_paths)
+
+# Forward pass an input tensor through the model
+x = torch.randn(1, 3, 224, 224).to(device)
+feature_outputs, output = hooked_model(x)
+
+# Print the output shapes
+for layer_path, feature_output in feature_outputs.items():
+    print(f'{layer_path} output shape: {feature_output.shape}')
diff --git a/requirements.txt b/requirements.txt
index 204c1e7..d106dbd 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile pyproject.toml
+#    uv pip compile pyproject.toml -o requirements.txt
 filelock==3.13.4
     # via torch
 fsspec==2024.3.1