pytorch hook

1 files changed, 17 insertions, 0 deletions

diff --git a/cv/pretrained/features.py b/cv/pretrained/features.py
new file mode 100644
index 0000000..f479126
--- /dev/null
+++ b/cv/pretrained/features.py
@@ -0,0 +1,17 @@
+
+import timm
+import torch
+from torch import nn
+
+
+model_name = 'xception41'
+# model_name = 'resnet18'
+model = timm.create_model(model_name, pretrained=True)
+
+input = torch.randn(2, 3, 299, 299)
+
+o1 = model(input)
+print(o1.shape)
+
+o2 = model.forward_features(input)
+print(o2.shape)
\ No newline at end of file