torch 网络参数查看,参数量统计

以LeNet为例：

import torch                                                                        
from torchsummary import summary                                                    
                                                                                    
class LeNet5(torch.nn.Module):                                                      
    def __init__(self, num_classes):                                                
        super(LeNet5, self).__init__()                                              
        self.features = torch.nn.Sequential(                                        
            torch.nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5),         
            torch.nn.ReLU(),                                                        
            torch.nn.MaxPool2d(kernel_size=2, stride=2),                            
            torch.nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(3,5)),    
            torch.nn.ReLU(),                                                        
            torch.nn.MaxPool2d(kernel_size=2, stride=2),                            
            torch.nn.Conv2d(in_channels=64, out_channels=128, kernel_size=(1,7)),   
            torch.nn.ReLU(),                                                        
            torch.nn.MaxPool2d(kernel_size=2, stride=2),                            
            torch.nn.Conv2d(in_channels=128, out_channels=256, kernel_size=(1,5)),  
            torch.nn.ReLU(),                                                        
            torch.nn.MaxPool2d(kernel_size=1, stride=2),                            
        )                                                                           
        self.classifier = torch.nn.Sequential(                                      
            torch.nn.Linear(in_features=256, out_features=84),                      
            torch.nn.ReLU(),                                                        
            torch.nn.Linear(in_features=84, out_features=num_classes)               
        )                                                                           
                                                                                    
    def forward(self,x):                                                            
        x = self.features(x)                                                        
        x = x.view(x.size(0), -1)                                                   
        x = self.classifier(x)                                                      
        return x                                                                    
                                                                                    
if __name__ == "__main__":                                                                                                                             
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")                                                             
    model = LeNet5(num_classes=3).to(device)    #更改为自己网络加载                                    
    summary(model, (3, 32, 32))         # input c,h,w