卷积神经网络参数量计算及显存占用 |
您所在的位置：网站首页 › gpu占用和显存占用的区别 › 卷积神经网络参数量计算及显存占用 |
卷积神经网络参数量计算及显存占用
|
参数量的计算 卷积神经网络的参数位于卷积层和全连接层，其中一个全连接层要比卷积层多好多参数。 显存占用 显存占用主要是模型(卷积层和全连接层的参数所占的显存)和各层的输出所占的显存。 下面以VGG16为例来讲解。 总共参数量 138,344,128(138357544) 占用显存 参数占用显存 138344128 *4 /1024 /1024 = 527.74MB ~ 528MB(一共有138344128个参数，每个都是float32类型的，即一个占用32位=4bytes) 模型占用显存 15,237,608 *4 /1024 /1024 = 58.12MB / image
验证代码 # bias = False import torch import torch.nn as nn class VGG(nn.Module): def __init__(self, vgg_name, cfg, num_classes=10, bn=False): super(VGG, self).__init__() self.vgg_base = self.make_layer(cfg, bn) if vgg_name == 'vgg16_C': self.fc1 = nn.Sequential(nn.Linear(512 * 8 * 8, 4096), nn.ReLU(inplace=True), nn.Dropout()) else: self.fc1 = nn.Sequential(nn.Linear(512 * 7 * 7, 4096, bias= False), # 这里的4096只是一个经验值，当然可以变为别的数，但不要小于要预测的类别数 nn.ReLU(inplace=True), nn.Dropout()) self.fc2 = nn.Sequential(nn.Linear(4096, 4096,bias= False), nn.ReLU(inplace=True), nn.Dropout()) self.fc3 = nn.Linear(4096, num_classes,bias= False) def make_layer(self, cfg, bn=False): layers = [] in_channels = 3 for v in cfg: if v == 'M': layers += [nn.MaxPool2d((2, 2), stride=2)] else: out_channels, s = v.strip().split('_') # 例如512 3 out_channels, s = int(out_channels), int(s) if bn: layers += [nn.Conv2d(in_channels, out_channels, (s, s), padding=1, bias= False), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True)] # nn.Conv2d bias默认是true else: layers += [nn.Conv2d(in_channels, out_channels, (s, s), padding=1, bias= False), nn.ReLU(inplace=True)] in_channels = out_channels return nn.Sequential(*layers) def forward(self, x): batch_size = x.size()[0] x = self.vgg_base(x) print (x.shape) # [1,512,7,7] x = x.view(batch_size, -1) x = self.fc1(x) x = self.fc2(x) x = self.fc3(x) return x cfg = { 'vgg11_A': ['64_3', 'M', '128_3', 'M', '256_3', '256_3', 'M', '512_3', '512_3', 'M', '512_3', '512_3', 'M'], 'vgg13_B': ['64_3', '64_3', 'M', '128_3', '128_3', 'M', '256_3', '256_3', 'M', '512_3', '512_3', 'M', '512_3', '512_3', 'M'], 'vgg16_C': ['64_3', '64_3', 'M', '128_3', '128_3', 'M', '256_3', '256_3', '256_1', 'M', '512_3', '512_3', '512_1', 'M', '512_3', '512_3', '512_1', 'M'], 'vgg16_D': ['64_3', '64_3', 'M', '128_3', '128_3', 'M', '256_3', '256_3', '256_3', 'M', '512_3', '512_3', '512_3', 'M', '512_3', '512_3', '512_3', 'M'], 'vgg19_E': ['64_3', '64_3', 'M', '128_3', '128_3', 'M', '256_3', '256_3', '256_3', '256_3', 'M', '512_3', '512_3', '512_3', '512_3', 'M', '512_3', '512_3', '512_3', '512_3', 'M'], } if __name__ == '__main__': input_tensor = torch.randn((1, 3, 224, 224)) input_var = torch.autograd.Variable(input_tensor) vgg_name = 'vgg16_D' model = VGG(vgg_name, cfg[vgg_name], num_classes=1000) output = model(input_var) print (output.shape) print(output.dtype) print('total_parameter_num:', sum(param.numel() for param in model.parameters())) # bias=True import torch import torch.nn as nn class VGG(nn.Module): def __init__(self, vgg_name, cfg, num_classes=10, bn=False): super(VGG, self).__init__() self.vgg_base = self.make_layer(cfg, bn) if vgg_name == 'vgg16_C': self.fc1 = nn.Sequential(nn.Linear(512 * 8 * 8, 4096), nn.ReLU(inplace=True), nn.Dropout()) else: self.fc1 = nn.Sequential(nn.Linear(512 * 7 * 7, 4096), # 这里的4096只是一个经验值，当然可以变为别的数，但不要小于要预测的类别数 nn.ReLU(inplace=True), nn.Dropout()) self.fc2 = nn.Sequential(nn.Linear(4096, 4096), nn.ReLU(inplace=True), nn.Dropout()) self.fc3 = nn.Linear(4096, num_classes) def make_layer(self, cfg, bn=False): layers = [] in_channels = 3 for v in cfg: if v == 'M': layers += [nn.MaxPool2d((2, 2), stride=2)] else: out_channels, s = v.strip().split('_') # 例如512 3 out_channels, s = int(out_channels), int(s) if bn: layers += [nn.Conv2d(in_channels, out_channels, (s, s), padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True)] # nn.Conv2d bias默认是true else: layers += [nn.Conv2d(in_channels, out_channels, (s, s), padding=1), nn.ReLU(inplace=True)] in_channels = out_channels return nn.Sequential(*layers) def forward(self, x): batch_size = x.size()[0] x = self.vgg_base(x) print (x.shape) # [1,512,7,7] x = x.view(batch_size, -1) x = self.fc1(x) x = self.fc2(x) x = self.fc3(x) return x cfg = { 'vgg11_A': ['64_3', 'M', '128_3', 'M', '256_3', '256_3', 'M', '512_3', '512_3', 'M', '512_3', '512_3', 'M'], 'vgg13_B': ['64_3', '64_3', 'M', '128_3', '128_3', 'M', '256_3', '256_3', 'M', '512_3', '512_3', 'M', '512_3', '512_3', 'M'], 'vgg16_C': ['64_3', '64_3', 'M', '128_3', '128_3', 'M', '256_3', '256_3', '256_1', 'M', '512_3', '512_3', '512_1', 'M', '512_3', '512_3', '512_1', 'M'], 'vgg16_D': ['64_3', '64_3', 'M', '128_3', '128_3', 'M', '256_3', '256_3', '256_3', 'M', '512_3', '512_3', '512_3', 'M', '512_3', '512_3', '512_3', 'M'], 'vgg19_E': ['64_3', '64_3', 'M', '128_3', '128_3', 'M', '256_3', '256_3', '256_3', '256_3', 'M', '512_3', '512_3', '512_3', '512_3', 'M', '512_3', '512_3', '512_3', '512_3', 'M'], } if __name__ == '__main__': input_tensor = torch.randn((1, 3, 224, 224)) input_var = torch.autograd.Variable(input_tensor) vgg_name = 'vgg16_D' model = VGG(vgg_name, cfg[vgg_name], num_classes=1000) output = model(input_var) print (output.shape) print(output.dtype) print('total_parameter_num:', sum(param.numel() for param in model.parameters()))参考 https://blog.csdn.net/qq_38262728/article/details/89813503 https://blog.csdn.net/qian99/article/details/79008053 |
各层的具体情况如下
128 * 58.12 MB * 2（这里乘以2，因为forward和backward）【本文地址】
今日新闻 |
推荐新闻 |