mirror of
https://github.com/lordmathis/CUDANet.git
synced 2025-11-06 09:44:28 +00:00
21 lines
519 B
Python
21 lines
519 B
Python
import torchvision
|
|
import torch
|
|
import sys
|
|
|
|
from torchsummary import summary
|
|
|
|
sys.path.append('../../tools') # Ugly hack
|
|
from utils import export_model_weights, print_model_parameters
|
|
|
|
if __name__ == "__main__":
|
|
alexnet = torchvision.models.alexnet(pretrained=True)
|
|
print_model_parameters(alexnet) # print layer names and number of parameters
|
|
export_model_weights(alexnet, 'alexnet_weights.bin')
|
|
print()
|
|
|
|
if torch.cuda.is_available():
|
|
alexnet.cuda()
|
|
|
|
summary(alexnet, (3, 227, 227))
|
|
|