查看模型参数(以AlexNet为例)
import torch
import torch. nn as nn
import torchvision
class AlexNet ( nn. Module) :
def __init__ ( self, num_classes= 1000 ) :
super ( AlexNet, self) . __init__( )
self. feature_extraction = nn. Sequential(
nn. Conv2d( in_channels= 3 , out_channels= 96 , kernel_size= 11 , stride= 4 , padding= 2 , bias= False ) ,
nn. ReLU( inplace= True ) ,
nn. MaxPool2d( kernel_size= 3 , stride= 2 , padding= 0 ) ,
nn. Conv2d( in_channels= 96 , out_channels= 192 , kernel_size= 5 , stride= 1 , padding= 2 , bias= False ) ,
nn. ReLU( inplace= True ) ,
nn. MaxPool2d( kernel_size= 3 , stride= 2 , padding= 0 ) ,
nn. Conv2d( in_channels= 192 , out_channels= 384 , kernel_size= 3 , stride= 1 , padding= 1 , bias= False ) ,
nn. ReLU( inplace= True ) ,
nn. Conv2d( in_channels= 384 , out_channels= 256 , kernel_size= 3 , stride= 1 , padding= 1 , bias= False ) ,
nn. ReLU( inplace= True ) ,
nn. Conv2d( in_channels= 256 , out_channels= 256 , kernel_size= 3 , stride= 1 , padding= 1 , bias= False ) ,
nn. ReLU( inplace= True ) ,
nn. MaxPool2d( kernel_size= 3 , stride= 2 , padding= 0 ) ,
)
self. classifier = nn. Sequential(
nn. Dropout( p= 0.5 ) ,
nn. Linear( in_features= 256 * 6 * 6 , out_features= 4096 ) ,
nn. ReLU( inplace= True ) ,
nn. Dropout( p= 0.5 ) ,
nn. Linear( in_features= 4096 , out_features= 4096 ) ,
nn. ReLU( inplace= True ) ,
nn. Linear( in_features= 4096 , out_features= num_classes) ,
)
def forward ( self, x) :
x = self. feature_extraction( x)
x = x. view( x. size( 0 ) , 256 * 6 * 6 )
x = self. classifier( x)
return x
if __name__ == '__main__' :
model = AlexNet( )
for name, parameters in model. named_parameters( ) :
print ( name, ':' , parameters. size( ) )
feature_extraction.0.weight : torch.Size([96, 3, 11, 11])
feature_extraction.3.weight : torch.Size([192, 96, 5, 5])
feature_extraction.6.weight : torch.Size([384, 192, 3, 3])
feature_extraction.8.weight : torch.Size([256, 384, 3, 3])
feature_extraction.10.weight : torch.Size([256, 256, 3, 3])
classifier.1.weight : torch.Size([4096, 9216])
classifier.1.bias : torch.Size([4096])
classifier.4.weight : torch.Size([4096, 4096])
classifier.4.bias : torch.Size([4096])
classifier.6.weight : torch.Size([1000, 4096])
classifier.6.bias : torch.Size([1000])
计算参数量与可训练参数量
def get_parameter_number ( model) :
total_num = sum ( p. numel( ) for p in model. parameters( ) )
trainable_num = sum ( p. numel( ) for p in model. parameters( ) if p. requires_grad)
return { 'Total' : total_num, 'Trainable' : trainable_num}
第三方工具
from torchstat import stat
import torchvision. models as models
model = models. alexnet( )
stat( model, ( 3 , 224 , 224 ) )
from torchvision. models import alexnet
import torch
from thop import profile
model = alexnet( )
input = torch. randn( 1 , 3 , 224 , 224 )
flops, params = profile( model, inputs= ( input , ) )
print ( flops, params)