用PyTorch实现VGG16,包括数据集训练和测试的可视化,要求可以用自己的数据,网络参数可改2
时间: 2023-07-12 16:08:25 浏览: 100
好的,以下是用PyTorch实现VGG16的代码。我们将使用CIFAR-10数据集进行训练和测试,并且我们将添加一些数据增强技术以提高模型的性能。
首先,我们需要导入必要的库和数据集:
```python
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
# 数据增强
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# 加载数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=transform_train)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128,
shuffle=True, num_workers=2)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
download=True, transform=transform_test)
testloader = torch.utils.data.DataLoader(testset, batch_size=128,
shuffle=False, num_workers=2)
classes = ('plane', 'car', 'bird', 'cat',
'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
```
接下来,我们将定义VGG16模型。由于我们可以改变网络参数,所以我们将定义一个函数来接受不同的参数:
```python
def VGG16(num_classes=10, num_channels=3):
layers = []
in_channels = num_channels
# 第一个卷积块
layers += [nn.Conv2d(in_channels, 64, kernel_size=3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True)]
layers += [nn.Conv2d(64, 64, kernel_size=3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True)]
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
# 第二个卷积块
layers += [nn.Conv2d(64, 128, kernel_size=3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True)]
layers += [nn.Conv2d(128, 128, kernel_size=3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True)]
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
# 第三个卷积块
layers += [nn.Conv2d(128, 256, kernel_size=3, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True)]
layers += [nn.Conv2d(256, 256, kernel_size=3, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True)]
layers += [nn.Conv2d(256, 256, kernel_size=3, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True)]
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
# 第四个卷积块
layers += [nn.Conv2d(256, 512, kernel_size=3, padding=1),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True)]
layers += [nn.Conv2d(512, 512, kernel_size=3, padding=1),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True)]
layers += [nn.Conv2d(512, 512, kernel_size=3, padding=1),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True)]
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
# 第五个卷积块
layers += [nn.Conv2d(512, 512, kernel_size=3, padding=1),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True)]
layers += [nn.Conv2d(512, 512, kernel_size=3, padding=1),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True)]
layers += [nn.Conv2d(512, 512, kernel_size=3, padding=1),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True)]
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
# 全连接层
layers += [nn.Flatten()]
layers += [nn.Linear(512 * 2 * 2, 4096),
nn.ReLU(inplace=True),
nn.Dropout()]
layers += [nn.Linear(4096, 4096),
nn.ReLU(inplace=True),
nn.Dropout()]
layers += [nn.Linear(4096, num_classes)]
return nn.Sequential(*layers)
```
接下来,我们将定义损失函数、优化器和学习率调度程序:
```python
device = 'cuda' if torch.cuda.is_available() else 'cpu'
net = VGG16(num_classes=10, num_channels=3)
net.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
lr_scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[100, 150], gamma=0.1)
```
现在我们可以开始训练模型了:
```python
def train(net, trainloader, optimizer, criterion, device):
net.train()
train_loss = 0
correct = 0
total = 0
for batch_idx, (inputs, targets) in enumerate(trainloader):
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
train_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
return train_loss / (batch_idx + 1), 100. * correct / total
def test(net, testloader, criterion, device):
net.eval()
test_loss = 0
correct = 0
total = 0
with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(testloader):
inputs, targets = inputs.to(device), targets.to(device)
outputs = net(inputs)
loss = criterion(outputs, targets)
test_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
return test_loss / (batch_idx + 1), 100. * correct / total
for epoch in range(200):
train_loss, train_acc = train(net, trainloader, optimizer, criterion, device)
test_loss, test_acc = test(net, testloader, criterion, device)
lr_scheduler.step()
print('Epoch %d: Train Loss: %.3f | Train Acc: %.3f%% | Test Loss: %.3f | Test Acc: %.3f%%'
% (epoch + 1, train_loss, train_acc, test_loss, test_acc))
```
最后,我们可以使用Matplotlib来可视化一些测试图像和它们的预测结果:
```python
import matplotlib.pyplot as plt
import numpy as np
def imshow(img):
img = img / 2 + 0.5 # unnormalize
npimg = img.numpy()
plt.imshow(np.transpose(npimg, (1, 2, 0)))
plt.show()
dataiter = iter(testloader)
images, labels = dataiter.next()
# 显示图像
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))
# 预测结果
outputs = net(images.to(device))
_, predicted = torch.max(outputs, 1)
print('Predicted: ', ' '.join('%5s' % classes[predicted[j]] for j in range(4)))
```
这就是用PyTorch实现VGG16的完整代码。你可以通过改变参数来改变网络结构或使用不同的数据集进行训练和测试。
阅读全文