Pytorch简单训练猫狗-快速上手

贴一个能直接运行的猫狗分类，方便新手上路
答辩完了再把我的项目放上来。

自己搭好环境，按顺序复制，运行就行。

猫狗分类数据集：链接：https://pan.baidu.com/s/1vcZ3VHNUVlR9y9WQurSu5w
提取码：s3ji

引用包

from torchvision import transforms
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
from torchvision.datasets import ImageFolder
from torch.utils.data import Dataset

处理图片，统一大小224，变换格式为Tensor()，最后一个搞不搞把，看他们都这么写，就跟着弄

simple_transform = transforms.Compose([transforms.Resize((224, 224))
                                          , transforms.ToTensor()
                                          , transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                                       ])

简单模型，
第一层3，10意思是输入一个三通道的图片，就是（R，G，B）三通道，就是彩色图片，然后用10个随机参数的卷积核去卷积它，所以输出也是10。卷积核大小为5

第二层10，20 就是再用20卷积核个去卷积.

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(56180, 500)
        self.fc2 = nn.Linear(500, 50)
        self.fc3 = nn.Linear(50, 2)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = F.relu(self.fc2(x))
        x = F.dropout(x, training=self.training)
        x = self.fc3(x)
        return F.log_softmax(x, dim=1)

训练。核心是这个几个：
optimizer.zero_grad() 梯度归零
loss.backward()#回归损失
optimizer.step()

其他的也不要忽视：
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
#优化方式，这是带动量的随机梯度，不填momentum就不带动量了，带上好。
F.nll_loss(output, target)#求损失
running_correct += preds.eq(target.data.view_as(preds)).cpu().sum()
这个是判断是否正确的代码

# 训练
def fit(epoch, model, data_loader, phase='训练', volatile=False, is_cuda=True):
    if phase == '训练':
        model.train()
    if phase == '验证':
        model.eval()
        volatile = True
    running_loss = 0.0
    running_correct = 0
    # 设置优化
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
    for batch_idx, (data, target) in enumerate(data_loader):
        if is_cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data, volatile), Variable(target)
        if phase == '训练':
            optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)

        running_loss += F.nll_loss(output, target, reduction='sum').item()
        preds = output.data.max(dim=1, keepdim=True)[1]
        running_correct += preds.eq(target.data.view_as(preds)).cpu().sum()
        if phase == '训练':
            loss.backward()
            optimizer.step()

    loss = running_loss / len(data_loader.dataset)
    accuracy = float(100. * running_correct / len(data_loader.dataset))
    if phase == '训练':
        print(f":第{epoch}代：平均损失 {loss:{5}.{2}} 正确率 {accuracy:{7}.{4}}%", end=" ")
    else:
        print(f"验证损失 {loss:{5}.{2}} ,验证正确率 {accuracy:{7}.{4}}%")

在这加载数据，设置训练前需要的数据

def train():
    # 训练损失，准确度
    path_root = 'F:/DogVsCat/'  # 数据集的位置
    model_name = "model-1-20.pth"
    is_cuda = False  # 检查是否用GPU
    if torch.cuda.is_available():
        is_cuda = True
    # 加载数据
    train_data = ImageFolder(path_root + 'train/', simple_transform)
    valid_data = ImageFolder(path_root + 'valid/', simple_transform)
    # 加载训练和测试数据
    train_data_loader = torch.utils.data.DataLoader(train_data, batch_size=32,shuffle=True)
    valid_data_loader = torch.utils.data.DataLoader(valid_data, batch_size=32,shuffle=True)
    model = Net()
    if is_cuda:
        model.cuda()
    print(f'数据集包含类别：{train_data.classes}')
    for epoch in range(1, 21):  # 20
        fit(epoch, model, train_data_loader, phase='训练')
        fit(epoch, model, valid_data_loader, phase='验证')
    torch.save(model.state_dict(), model_name)
    print("训练完毕，模型已保存为", model_name)

主函数运行

if __name__ == "__main__":
    train()