How to check the network parameter explicit memory usage by pytorch
This article mainly introduces pytorch how to view the network parameters of apparent memory usage, the article introduces in great detail, has a certain reference value, interested friends must read it!
1. Use torchstatpip install torchstat from torchstat import statimport torchvision.models as modelsmodel = models.resnet152 () stat (model, (3,224,224))
With regard to the parameters of the stat function, the first should be the model, the second is the input size, and 3 is the number of channels. I didn't investigate the detailed parameters of the function, and I don't know why I don't prompt for the corresponding parameters when using it.
two。 Use torchsummarypip install torchsummary from torchsummary import summarysummary (model.cuda (), input_size= (3jin32), batch_size=-1)
Using this function to prompt the parameters directly, you can find that there is an explicit input batch_size directly, and I feel as if the function is better. But! For some reason, this function keeps reporting errors on my machine!
TypeError: can't convert CUDA tensor to numpy. Use Tensor.cpu () to copy the tensor to host memory first.
Update: after consultation on the forum, the reason for reporting the error has been found, just put the
Pip install torchsummary
Modify to
Pip install torch-summary
Add: Pytorch looks at the model parameters and calculates the number of model parameters and trainable parameters
View model parameters (take AlexNet as an example) import torchimport torch.nn as nnimport torchvisionclass 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 (pendant 0.5), nn.Linear (in_features=256*6*6) Out_features=4096), nn.ReLU (inplace=True), nn.Dropout (pendant 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 xif _ _ name__ ='_ _ main__': # model = torchvision.models.AlexNet () model = AlexNet () # print model parameter # for param in model.parameters (): # print (param) # print model name and shape 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]) calculates the number of parameters and the number of trainable parameters 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} third-party tools from torchstat import statimport torchvision.models as modelsmodel = models.alexnet () stat (model, (3,224,224)) from torchvision.models import alexnetimport torchfrom thop import profilemodel = alexnet () input = torch.randn (1,3,224,224) flops, params = profile (model, inputs= (input,)) print (flops, params) these are all the contents of the article "how to check the apparent memory usage of pytorch parameters" Thank you for reading! Hope to share the content to help you, more related knowledge, welcome to follow the industry information channel!
