# -*- coding:utf-8 -*-
"""
开发时间:2024/8/4 16:43
@author:xd
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
class AlexNet(nn.Module):
def __init__(self, num_classes):
super(AlexNet, self).__init__()
self.conv1 = nn.Conv2d(1, 96, kernel_size=11, stride=4)
self.max_pool1 = nn.MaxPool2d(kernel_size=3, stride=2)
self.conv2 = nn.Conv2d(96, 256, kernel_size=5, stride=1, padding=2)
self.max_pool2 = nn.MaxPool2d(kernel_size=3, stride=2)
self.conv3 = nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1)
self.conv4 = nn.Conv2d(384, 384, kernel_size=3, stride=1, padding=1)
self.conv5 = nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1)
self.max_pool3 = nn.MaxPool2d(kernel_size=3, stride=2)
self.fc6 = nn.Linear(256 * 6 * 6, 4096)
self.dropout1 = nn.Dropout(0.5)
self.fc7 = nn.Linear(4096, 4096)
self.dropout2 = nn.Dropout(0.5)
self.fc8 = nn.Linear(4096, num_classes)
self.softmax = nn.Softmax(dim=1)
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.max_pool1(x)
x = F.relu(self.conv2(x))
x = self.max_pool2(x)
x = F.relu(self.conv3(x))
x = F.relu(self.conv4(x))
x = F.relu(self.conv5(x))
x = self.max_pool3(x)
x = torch.flatten(x, 1)
x = F.relu(self.fc6(x))
x = self.dropout1(x)
x = F.relu(self.fc7(x))
x = self.dropout2(x)
x = self.softmax(self.fc8(x))
return x
def test():
model = AlexNet(num_classes=10)
print(model)
input = torch.zeros((1, 1, 227, 227))
output = model(input)
print(output)
if __name__ == "__main__":
test()
AlexNet代码复现
于 2024-08-04 18:17:43 首次发布