Python可以调用任意框架,完成对tiny-imagenet的训练过程和分类预测。
时间: 2023-06-19 09:08:49 浏览: 99
是的,Python可以调用各种深度学习框架,例如TensorFlow、PyTorch、Keras等,来完成对tiny-imagenet的训练过程和分类预测。以下是一个使用PyTorch完成tiny-imagenet训练和预测的示例代码:
1. 训练模型
```python
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision.transforms import transforms
# 定义数据预处理
data_transforms = transforms.Compose([
transforms.RandomResizedCrop(64),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
# 加载训练数据
train_dataset = ImageFolder('tiny-imagenet/train', transform=data_transforms)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
# 定义模型
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(64, 128, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(128, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(256, 512, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
)
self.fc = nn.Sequential(
nn.Linear(512 * 2 * 2, 1024),
nn.ReLU(inplace=True),
nn.Dropout(0.5),
nn.Linear(1024, 200),
)
def forward(self, x):
x = self.conv(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
# 定义损失函数和优化器
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Model().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# 训练模型
for epoch in range(50):
train_loss = 0.0
train_correct = 0
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()
train_loss += loss.item() * data.size(0)
_, pred = torch.max(output, 1)
train_correct += (pred == target).sum().item()
train_loss /= len(train_loader.dataset)
train_acc = train_correct / len(train_loader.dataset)
print('Epoch: {}, Train Loss: {:.4f}, Train Acc: {:.4f}'.format(epoch+1, train_loss, train_acc))
```
2. 预测
```python
from PIL import Image
# 加载测试图片
img_path = 'tiny-imagenet/test/images/test_1.JPEG'
img = Image.open(img_path)
img = data_transforms(img).unsqueeze(0)
# 加载模型
model = Model().to(device)
model.load_state_dict(torch.load('model.pth'))
# 预测图片类别
model.eval()
img = img.to(device)
output = model(img)
_, pred = torch.max(output, 1)
print('Predicted: {}'.format(pred.item()))
```
需要注意的是,以上示例代码仅供参考,实际应用中还需要根据具体问题进行相应的调整和优化。
阅读全文