请你扮演一个程序员,请你在pytorch框架下,比较用同一个vgg16模型,在同样epoch训练轮数下,比较不同优化器、不同学习率更新方法下的ROC曲线,请写出实际代码。
时间: 2023-06-14 22:04:16 浏览: 81
pytorch迁移学习训练VGG16和模型测试代码
好的,下面是代码实现。首先需要导入必要的库:PyTorch、NumPy、Scikit-learn、Matplotlib。
```python
import torch
import numpy as np
from sklearn import metrics
import matplotlib.pyplot as plt
```
接下来定义 VGG16 模型:
```python
class VGG16(torch.nn.Module):
def __init__(self):
super(VGG16, self).__init__()
self.features = torch.nn.Sequential(
torch.nn.Conv2d(3, 64, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
torch.nn.Conv2d(64, 64, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
torch.nn.MaxPool2d(kernel_size=2, stride=2),
torch.nn.Conv2d(64, 128, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
torch.nn.Conv2d(128, 128, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
torch.nn.MaxPool2d(kernel_size=2, stride=2),
torch.nn.Conv2d(128, 256, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
torch.nn.MaxPool2d(kernel_size=2, stride=2),
torch.nn.Conv2d(256, 512, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
torch.nn.Conv2d(512, 512, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
torch.nn.Conv2d(512, 512, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
torch.nn.MaxPool2d(kernel_size=2, stride=2),
torch.nn.Conv2d(512, 512, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
torch.nn.Conv2d(512, 512, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
torch.nn.Conv2d(512, 512, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
torch.nn.MaxPool2d(kernel_size=2, stride=2),
)
self.classifier = torch.nn.Sequential(
torch.nn.Linear(512 * 7 * 7, 4096),
torch.nn.ReLU(inplace=True),
torch.nn.Dropout(),
torch.nn.Linear(4096, 4096),
torch.nn.ReLU(inplace=True),
torch.nn.Dropout(),
torch.nn.Linear(4096, 2),
)
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), -1)
x = self.classifier(x)
return x
```
接下来分别定义使用不同优化器和学习率更新方法下的训练代码:
```python
def train_SGD(model, train_loader, optimizer, criterion, device):
model.train()
for data, target in train_loader:
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
def train_RMSprop(model, train_loader, optimizer, criterion, device):
model.train()
for data, target in train_loader:
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
def train_Adam(model, train_loader, optimizer, criterion, device):
model.train()
for data, target in train_loader:
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
```
定义测试函数:
```python
def test(model, test_loader, device):
model.eval()
y_true = []
y_pred = []
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
y_true.extend(target.cpu().numpy())
y_pred.extend(torch.argmax(output, dim=1).cpu().numpy())
fpr, tpr, thresholds = metrics.roc_curve(y_true, y_pred)
roc_auc = metrics.auc(fpr, tpr)
return fpr, tpr, roc_auc
```
最后定义主函数:
```python
def main():
# 超参数
lr = 0.001
epochs = 10
batch_size = 32
# 加载数据集
train_data = ...
test_data = ...
train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=False)
# 定义设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 定义模型
model = VGG16().to(device)
# 定义损失函数
criterion = torch.nn.CrossEntropyLoss()
# 使用 SGD 优化器训练模型并测试
optimizer_SGD = torch.optim.SGD(model.parameters(), lr=lr)
fpr_SGD, tpr_SGD, roc_auc_SGD = [], [], []
for epoch in range(epochs):
train_SGD(model, train_loader, optimizer_SGD, criterion, device)
fpr, tpr, roc_auc = test(model, test_loader, device)
fpr_SGD.append(fpr)
tpr_SGD.append(tpr)
roc_auc_SGD.append(roc_auc)
# 使用 RMSprop 优化器训练模型并测试
optimizer_RMSprop = torch.optim.RMSprop(model.parameters(), lr=lr)
fpr_RMSprop, tpr_RMSprop, roc_auc_RMSprop = [], [], []
for epoch in range(epochs):
train_RMSprop(model, train_loader, optimizer_RMSprop, criterion, device)
fpr, tpr, roc_auc = test(model, test_loader, device)
fpr_RMSprop.append(fpr)
tpr_RMSprop.append(tpr)
roc_auc_RMSprop.append(roc_auc)
# 使用 Adam 优化器训练模型并测试
optimizer_Adam = torch.optim.Adam(model.parameters(), lr=lr)
fpr_Adam, tpr_Adam, roc_auc_Adam = [], [], []
for epoch in range(epochs):
train_Adam(model, train_loader, optimizer_Adam, criterion, device)
fpr, tpr, roc_auc = test(model, test_loader, device)
fpr_Adam.append(fpr)
tpr_Adam.append(tpr)
roc_auc_Adam.append(roc_auc)
# 绘制 ROC 曲线
plt.figure()
plt.plot(fpr_SGD, tpr_SGD, label="SGD (AUC = %.4f)" % roc_auc_SGD[-1])
plt.plot(fpr_RMSprop, tpr_RMSprop, label="RMSprop (AUC = %.4f)" % roc_auc_RMSprop[-1])
plt.plot(fpr_Adam, tpr_Adam, label="Adam (AUC = %.4f)" % roc_auc_Adam[-1])
plt.plot([0, 1], [0, 1], "k--")
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.0])
plt.xlabel("False Positive Rate")
plt.ylabel("True Positive Rate")
plt.title("ROC Curve")
plt.legend(loc="lower right")
plt.show()
if __name__ == "__main__":
main()
```
这份代码中,我们使用了三种不同的优化器:SGD,RMSprop 和 Adam。在每个 epoch 中,我们都会使用训练集对模型进行训练,并使用测试集对模型进行测试,得到 ROC 曲线和 AUC 值。最后,我们绘制出三种优化器的 ROC 曲线,以便比较它们的效果。
阅读全文