Pytorch中怎么使用vis.line 来画train_loss train_acc test_acc 这三条线在同一个图中 并给出一个在深度学习中以网络为resnet50 、数据集为cifar10、把输出的模型保存、输出的数据保存到csv文件中、把代码模块化的实战代码
时间: 2024-05-08 13:16:00 浏览: 96
首先,我们需要安装visdom库,可以使用以下命令:
```python
!pip install visdom
```
接下来,需要在代码中引入相关库和模块:
```python
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import visdom
import csv
```
然后,定义一个函数来加载数据集:
```python
def load_dataset():
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')
return trainloader, testloader
```
接下来,定义一个函数来构建ResNet50网络:
```python
class ResNet50(nn.Module):
def __init__(self, num_classes=10):
super(ResNet50, self).__init__()
self.resnet50 = torchvision.models.resnet50(pretrained=False)
num_ftrs = self.resnet50.fc.in_features
self.resnet50.fc = nn.Linear(num_ftrs, num_classes)
def forward(self, x):
x = self.resnet50(x)
return x
```
然后,定义一个函数来训练模型:
```python
def train(model, trainloader, criterion, optimizer, epoch, device, vis):
model.train()
train_loss = 0.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 = model(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()
if batch_idx % 10 == 0:
vis.line(X=torch.Tensor([epoch+batch_idx/len(trainloader)]), Y=torch.Tensor([train_loss/len(trainloader)]), win='train_loss', update='append', name='train_loss')
vis.line(X=torch.Tensor([epoch+batch_idx/len(trainloader)]), Y=torch.Tensor([100.*correct/total]), win='train_acc', update='append', name='train_acc')
print('Train Epoch: {} Loss: {:.6f} Acc: {:.6f}'.format(
epoch, train_loss/len(trainloader), 100.*correct/total))
```
接下来,定义一个函数来测试模型:
```python
def test(model, testloader, criterion, epoch, device, vis):
model.eval()
test_loss = 0.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 = model(inputs)
loss = criterion(outputs, targets)
test_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
if batch_idx % 10 == 0:
vis.line(X=torch.Tensor([epoch+batch_idx/len(testloader)]), Y=torch.Tensor([100.*correct/total]), win='test_acc', update='append', name='test_acc')
print('Test Epoch: {} Loss: {:.6f} Acc: {:.6f}'.format(
epoch, test_loss/len(testloader), 100.*correct/total))
```
接下来,定义主函数来调用以上函数:
```python
def main():
vis = visdom.Visdom()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
trainloader, testloader = load_dataset()
model = ResNet50().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
for epoch in range(100):
train(model, trainloader, criterion, optimizer, epoch, device, vis)
test(model, testloader, criterion, epoch, device, vis)
torch.save(model.state_dict(), 'resnet50_cifar10.pth')
with open('output.csv', mode='w') as file:
writer = csv.writer(file)
writer.writerow(['epoch', 'train_loss', 'train_acc', 'test_acc'])
for epoch in range(100):
train_loss = vis.line(Y=torch.Tensor([0]), win='train_loss', name='train_loss', update='none')
train_acc = vis.line(Y=torch.Tensor([0]), win='train_acc', name='train_acc', update='none')
test_acc = vis.line(Y=torch.Tensor([0]), win='test_acc', name='test_acc', update='none')
writer.writerow([epoch+1, train_loss['Y'][-1], train_acc['Y'][-1], test_acc['Y'][-1]])
```
最后,运行主函数即可进行训练、测试、模型保存和输出数据到csv文件的操作:
```python
if __name__ == '__main__':
main()
```
完整代码如下: