PyTorch实现ResNet34网络的完整攻略
ResNet是深度学习中非常流行的卷积神经网络之一,它在ImageNet数据集上取了非常好的效果。本文将详细讲解如何使用PyTorch实现ResNet34网络,包括数据预处理、网络搭建、训练和测试等内容。
数据预处理
在使用PyTorch实现ResNet34网络之前,需要对数据进行预处理。可以按照以下步骤进行预处理:
- 将图像缩放到256×256大小。
- 随机裁剪224×224大小的图像。
- 将图像水平翻转。
- 将像素值标准化为[0, 1]。
可以以下代码实现数据预处理:
import torchvision.transforms as transforms
transform_train = transforms.Compose([
transforms.Resize(256),
transforms.RandomCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
transform_test = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
网络搭建
在数据预处理完成后,需要搭建ResNet34网络。可以使用以下代码实现网络搭建:
import torch.nn as nn
import torch.nn.functional as F
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, in_planes, planes, stride=1):
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.shortcut = nn.Sequential()
if stride != 1 or in_planes != self.expansion*planes:
self.shortcut = nn.Sequential(
nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(self.expansion*planes)
)
def forward(self, x):
out = F.relu(self.bn1(self.conv1(x)))
out = self.bn2(self.conv2(out))
out += self.shortcut(x)
out = F.relu(out)
return out
class ResNet(nn.Module):
def __init__(self, block, num_blocks, num_classes=10):
super(ResNet, self).__init__()
self.in_planes = 64
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
self.linear = nn.Linear(512*block.expansion, num_classes)
def _make_layer(self, block, planes, num_blocks, stride):
strides = [stride] + [1]*(num_blocks-1)
layers = []
for stride in strides:
layers.append(block(self.in_planes, planes, stride))
self.in_planes = planes * block.expansion
return nn.Sequential(*layers)
def forward(self, x):
out = F.relu(self.bn1(self.conv1(x)))
out = self.layer1(out)
out = self.layer2(out)
out = self.layer3(out)
out = self.layer4(out)
out = F.avg_pool2d(out, 4)
out = out.view(out.size(0), -1)
out = self.linear(out)
return out
def ResNet34():
return ResNet(BasicBlock, [3,4,6,3])
在上面的代码中,BasicBlock
是ResNet中的基本块,ResNet
是ResNet网络,ResNet34
是ResNet34网络。
训练和测试
在网络搭建完成后,需要进行训练和测试。可以使用以下代码实现训练和测试:
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.models as models
from torch.utils.data import DataLoader
# 加载数据集
trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
trainloader = DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
testset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
testloader = DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)
# 定义模型、损失函数和优化器
net = ResNet34()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
# 训练模型
for epoch in range(100):
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 100 == 99:
print('[%d, %5d] loss: %.3f' % (epoch+1, i+1, running_loss/100))
running_loss = 0.0
# 测试模型
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))
在上面的代码中,首先加载CIFAR10数据集,然后定义ResNet34网络、交叉熵损失函数和随机梯度下降优化器。接着,使用训练集对模型进行训练,并使用测试集对模型进行测试。
示例1:使用ResNet34网络进行图像分类
以下是使用ResNet34网络进行图像分类的示例:
import torch
import torchvision.transforms as transforms
from PIL import Image
# 加载模型
net = ResNet34()
net.load_state_dict(torch.load('resnet34.pth'))
# 加载图像
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
image = Image.open('test.jpg')
image = transform(image)
image = image.unsqueeze(0)
# 预测图像类别
outputs = net(image)
_, predicted = torch.max(outputs.data, 1)
print(predicted.item())
在上面的示例中,首先加载ResNet34网络和预训练模型,然后加载测试图像并进行预处理,最后使用ResNet34网络对图像进行分类。
示例2:使用ResNet34网络进行迁移学习
以下是使用ResNet34网络进行迁移学习的示例:
import torch
import torchvision.transforms as transforms
import torchvision.models as models
from torch.utils.data import DataLoader
from torch import nn, optim
# 加载数据集
trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
trainloader = DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
testset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
testloader = DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)
# 加载预训练模型
resnet = models.resnet34(pretrained=True)
for param in resnet.parameters():
param.requires_grad = False
# 替换最后一层全连接层
num_ftrs = resnet.fc.in_features
resnet.fc = nn.Linear(num_ftrs, 10)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(resnet.fc.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
# 训练模型
for epoch in range(100):
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data
optimizer.zero_grad()
outputs = resnet(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 100 == 99:
print('[%d, %5d] loss: %.3f' % (epoch+1, i+1, running_loss/100))
running_loss = 0.0
# 测试模型
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
outputs = resnet(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))
在上面的示例中,首先加载CIFAR10数据集,然后加载ResNet34预训练模型并替换最后一层全连接层。接着,使用训练集对模型进行训练,并使用测试集对模型进行测试。