输出利用resnet18预训练cifar10并保存训练权重的python代码
时间: 2023-12-21 08:04:42 浏览: 86
以下是利用resnet18预训练cifar10并保存训练权重的Python代码:
```python
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
# Define transforms for training data
train_transform = 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))
])
# Define transforms for test data
test_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# Load CIFAR10 dataset
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=train_transform)
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=test_transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=128,
shuffle=False, num_workers=2)
# Define ResNet18 model
net = torchvision.models.resnet18(pretrained=False, num_classes=10)
# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# Train the model
for epoch in range(10): # loop over the dataset multiple times
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % 100 == 99: # print every 100 mini-batches
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 100))
running_loss = 0.0
print('Finished Training')
# Save the trained model weights
PATH = './resnet18_cifar10.pth'
torch.save(net.state_dict(), PATH)
```
在上面的代码中,我们首先定义了训练和测试数据的transforms,并利用`torchvision.datasets.CIFAR10`加载CIFAR10数据集。然后,我们定义了ResNet18模型,以及损失函数和优化器。我们使用SGD优化器,学习率设置为0.001,动量设置为0.9。接下来,我们训练了模型,共进行了10个epochs,每个epoch中使用128个mini-batch。在训练过程中,我们输出了每100个mini-batches的平均损失。最后,我们保存了训练好的模型权重到文件中。