Mathis/CUDANet
1
0
Fork 0
mirror of https://github.com/lordmathis/CUDANet.git synced 2025-11-05 17:34:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update alexnet preprocessing

Browse Source
This commit is contained in:
Matúš Námešný
2024-09-04 21:41:26 +02:00
parent 7b8c4bd811
commit e7ec6c91f8
2 changed files with 31 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
examples/alexnet/alexnet.py
View File

@@ -3,11 +3,17 @@ import sys
import torchvision
sys.path.append('../../tools') # Ugly hack
from utils import export_model_weights, print_model_parameters
from utils import export_model_weights, print_model_parameters, predict
if __name__ == "__main__":
alexnet = torchvision.models.alexnet(weights=torchvision.models.AlexNet_Weights.DEFAULT)
print_model_parameters(alexnet) # print layer names and number of parameters
export_model_weights(alexnet, 'alexnet_weights.bin')
# predict(alexnet, 'cat.jpg')
weights = torchvision.models.AlexNet_Weights.DEFAULT
alexnet = torchvision.models.alexnet(weights=weights)
# print_model_parameters(alexnet) # print layer names and number of parameters
export_model_weights(alexnet, 'alexnet_weights.bin')
# class_labels = weights.meta["categories"]
# prediction = predict(alexnet, "margot.jpg")
# print(prediction, class_labels[prediction])