首先需要下载tiny-imagenet数据集，并解压缩到指定目录。可以从官网下载：http://cs231n.stanford.edu/tiny-imagenet-200.zip
接下来，我们使用PyTorch框架进行训练和分类预测。
训练代码：
```python
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.models as models
import torchvision.transforms as transforms
import torchvision.datasets as datasets
# 定义数据预处理
train_transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomCrop(64, padding=4),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
val_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
# 加载数据集
train_dataset = datasets.ImageFolder('tiny-imagenet-200/train', train_transform)
val_dataset = datasets.ImageFolder('tiny-imagenet-200/val', val_transform)
# 定义数据加载器
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=128, shuffle=False, num_workers=4)
# 定义模型
model = models.resnet18(pretrained=True)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 200)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=0.0001)
# 训练模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
for epoch in range(10):
model.train()
for i, (inputs, targets) in enumerate(train_loader):
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
print('Epoch [%d/%d], Iter [%d/%d], Loss: %.4f' % (epoch+1, 10, i+1, len(train_loader), loss.item()))
model.eval()
correct = 0
total = 0
with torch.no_grad():
for inputs, targets in val_loader:
inputs, targets = inputs.to(device), targets.to(device)
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1)
total += targets.size(0)
correct += (predicted == targets).sum().item()
print('Epoch [%d/%d], Validation Accuracy: %.4f' % (epoch+1, 10, correct/total))
# 保存模型
torch.save(model.state_dict(), 'tiny-imagenet-200-resnet18.pth')
```
分类预测代码：
```python
import torch
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
# 定义数据预处理
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
# 加载测试数据集
test_dataset = datasets.ImageFolder('tiny-imagenet-200/test', transform)
# 定义数据加载器
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=4)
# 加载模型
model = models.resnet18(pretrained=False)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 200)
model.load_state_dict(torch.load('tiny-imagenet-200-resnet18.pth'))
# 预测
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
model.eval()
with torch.no_grad():
for inputs, _ in test_loader:
inputs = inputs.to(device)
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1)
print(predicted)
```
注意：由于tiny-imagenet数据集较大，训练过程可能需要较长时间。建议使用GPU加速训练过程。
最低0.47元/天 解锁文章
1374
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
