PyTorch基础-猫狗分类实战-10

训练模型并保存

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets,transforms,models
from torch.utils.data import Dataset
import sys

# 数据预处理
transform = transforms.Compose([
    transforms.RandomResizedCrop(224),# 对图像进行随机裁剪
    transforms.RandomRotation(20),# 随机旋转角度
    transforms.RandomHorizontalFlip(p=0.5),# 随机水平翻转
    transforms.ToTensor()# 变成tensor格式
]) # 数据增强

# 读取数据
root = "image"
train_dataset = datasets.ImageFolder(root + "/train",transform)
test_dataset = datasets.ImageFolder(root + "/test",transform)

# 导入数据
train_loader = torch.utils.data.DataLoader(train_dataset,batch_size=8,shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset,batch_size=8,shuffle=True)

classes = train_dataset.classes
classes_index = train_dataset.class_to_idx
print