使用流程
demo
建议查看官方demo -->
github
pip install torchsummary导入
from torchsummary import summary使用
# 参数说明 summary(your_model, input_size=(channels, H, W))
myNet = NET() #NET为自己定义的网络模型 data = [(3, 100, 100), (3, 100, 100), (3, 100, 100)] #多输入数据 summary(myNet, data, batch_size=1, device='cpu')官方说明 demo
# CNN import torch import torch.nn as nn import torch.nn.functional as F from torchsummary import summary class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 10, kernel_size=5) self.conv2 = nn.Conv2d(10, 20, kernel_size=5) self.conv2_drop = nn.Dropout2d() self.fc1 = nn.Linear(320, 50) self.fc2 = nn.Linear(50, 10) def forward(self, x): x = F.relu(F.max_pool2d(self.conv1(x), 2)) x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) x = x.view(-1, 320) x = F.relu(self.fc1(x)) x = F.dropout(x, training=self.training) x = self.fc2(x) return F.log_softmax(x, dim=1) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # PyTorch v0.4.0 model = Net().to(device) summary(model, (1, 28, 28))
#VGG16 import torch from torchvision import models from torchsummary import summary device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') vgg = models.vgg16().to(device) summary(vgg, (3, 224, 224))
# 多输入 import torch import torch.nn as nn from torchsummary import summary class SimpleConv(nn.Module): def __init__(self): super(SimpleConv, self).__init__() self.features = nn.Sequential( nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1), nn.ReLU(), ) def forward(self, x, y): x1 = self.features(x) x2 = self.features(y) return x1, x2 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = SimpleConv().to(device) summary(model, [(1, 16, 16), (1, 28, 28)])
欢迎分享,转载请注明来源:内存溢出
评论列表(0条)