以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