GoogleNet(也称为 Inception v1)是由 Google 团队提出的一种深度卷积神经网络架构,以其高效的 Inception 模块和多层分类器著称。以下是使用 PyTorch 实现 GoogleNet 进行图像分类的步骤。
首先,确保你已经安装了必要的库:
pip install torch torchvision
注意:具体需要依据cuda版本来选择对应版本
2. 导入库
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
3. 定义 GoogleNet 模型
GoogleNet 的核心是 Inception 模块,它通过并行使用不同大小的卷积核来提取多尺度特征。
class InceptionModule(nn.Module):
def __init__(self, in_channels, out_1x1, reduce_3x3, out_3x3, reduce_5x5, out_5x5, out_pool):
super(InceptionModule, self).__init__()
# 1x1 卷积分支
self.branch1 = nn.Sequential(
nn.Conv2d(in_channels, out_1x1, kernel_size=1),
nn.BatchNorm2d(out_1x1),
nn.ReLU(inplace=True),
)
# 1x1 卷积 + 3x3 卷积分支
self.branch2 = nn.Sequential(
nn.Conv2d(in_channels, reduce_3x3, kernel_size=1),
nn.BatchNorm2d(reduce_3x3),
nn.ReLU(inplace=True),
nn.Conv2d(reduce_3x3, out_3x3, kernel_size=3, padding=1),
nn.BatchNorm2d(out_3x3),
nn.ReLU(inplace=True),
)
# 1x1 卷积 + 5x5 卷积分支
self.branch3 = nn.Sequential(
nn.Conv2d(in_channels, reduce_5x5, kernel_size=1),
nn.BatchNorm2d(reduce_5x5),
nn.ReLU(inplace=True),
nn.Conv2d(reduce_5x5, out_5x5, kernel_size=5, padding=2),
nn.BatchNorm2d(out_5x5),
nn.ReLU(inplace=True),
)
# 3x3 最大池化 + 1x1 卷积分支
self.branch4 = nn.Sequential(
nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
nn.Conv2d(in_channels, out_pool, kernel_size=1),
nn.BatchNorm2d(out_pool),
nn.ReLU(inplace=True),
)
def forward(self, x):
branch1 = self.branch1(x)
branch2 = self.branch2(x)
branch3 = self.branch3(x)
branch4 = self.branch4(x)
return torch.cat([branch1, branch2, branch3, branch4], 1)
class GoogleNet(nn.Module):
def __init__(self, num_classes=1000):
super(GoogleNet, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
)
self.conv2 = nn.Sequential(
nn.Conv2d(64, 64, kernel_size=1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.Conv2d(64, 192, kernel_size=3, padding=1),
nn.BatchNorm2d(192),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
)
self.inception3a = InceptionModule(192, 64, 96, 128, 16, 32, 32)
self.inception3b = InceptionModule(256, 128, 128, 192, 32, 96, 64)
self.maxpool3 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.inception4a = InceptionModule(480, 192, 96, 208, 16, 48, 64)
self.inception4b = InceptionModule(512, 160, 112, 224, 24, 64, 64)
self.inception4c = InceptionModule(512, 128, 128, 256, 24, 64, 64)
self.inception4d = InceptionModule(512, 112, 144, 288, 32, 64, 64)
self.inception4e = InceptionModule(528, 256, 160, 320, 32, 128, 128)
self.maxpool4 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.inception5a = InceptionModule(832, 256, 160, 320, 32, 128, 128)
self.inception5b = InceptionModule(832, 384, 192, 384, 48, 128, 128)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.dropout = nn.Dropout(0.4)
self.fc = nn.Linear(1024, num_classes)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.inception3a(x)
x = self.inception3b(x)
x = self.maxpool3(x)
x = self.inception4a(x)
x = self.inception4b(x)
x = self.inception4c(x)
x = self.inception4d(x)
x = self.inception4e(x)
x = self.maxpool4(x)
x = self.inception5a(x)
x = self.inception5b(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.dropout(x)
x = self.fc(x)
return x
4. 数据预处理和加载
使用 CIFAR-10 数据集作为示例:
transform = transforms.Compose([
transforms.Resize((224, 224)), # GoogleNet 输入尺寸为 224x224
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
5. 训练模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = GoogleNet(num_classes=10).to(device) # CIFAR-10 有 10 个类别
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
for epoch in range(10): # 训练 10 个 epoch
model.train()
running_loss = 0.0
for i, (inputs, labels) in enumerate(train_loader):
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 100 == 99: # 每 100 个 batch 打印一次损失
print(f'Epoch [{epoch + 1}/{10}], Step [{i + 1}/{len(train_loader)}], Loss: {running_loss / 100:.4f}')
running_loss = 0.0
print('Finished Training')
6. 测试模型
model.eval()
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in test_loader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Accuracy of the model on the test images: {100 * correct / total:.2f}%')
7. 保存模型
torch.save(model.state_dict(), 'googlenet_cifar10.pth')
8. 加载模型
model = GoogleNet(num_classes=10).to(device)
model.load_state_dict(torch.load('googlenet_cifar10.pth'))
以上代码展示了如何使用 PyTorch 实现 GoogleNet 进行图像分类。GoogleNet 的 Inception 模块通过多尺度特征提取提高了模型的表达能力,适合处理复杂的图像分类任务。可以根据需要调整模型结构或超参数,以适应不同的任务和数据集。