pytorch 常用API，分类训练总结

1.安装

https://pytorch.org/ 官网，根据对应情况执行命令行，就完了

2.数据预处理

from PIL import Image #读图
from torchvision import transforms, datasets #导入图像数据处理，转换库
import os
import torch

def readImg(path):
    '''
    用于替代ImageFolder的默认读取图片函数，以读取单通道图片
    '''
    return Image.open(path)


def ImageDataset(args):
    data_transforms = { #数据增强obj
        'train': transforms.Compose([ #组合转换
            transforms.RandomCrop(RANDOMCROP_SIZE), #随机裁剪
            transforms.RandomHorizontalFlip(), #随机水平翻转
            transforms.ToTensor(), #转tensor 变float变tensor
            transforms.Normalize([0.5], [0.5]) #改均值 
        ]),
        'test': transforms.Compose([#测试时要和train对应，有些增强不需要
            transforms.CenterCrop(RANDOMCROP_SIZE),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5])
        ]),
    }
    image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), #实用读分类图库
                        data_transforms[x] #数据增强
                        # , loader=readImg #单通道图像使用
                        )
                        for x in ['train', 'test']}
    dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x],
                    batch_size=args.batch_size, shuffle=(x == 'train'),
                    num_workers=args.num_workers)
                    for x in ['train', 'test']}
    dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'test']}
    class_names = image_datasets['train'].classes

    return dataloaders, dataset_sizes, class_names

3.网络模型保存和加载

(1) 仅仅保存和加载模型参数
torch.save(the_model.state_dict(), PATH)

the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))

(2) 保存和加载整个模型
torch.save(the_model, PATH)
#切记！ 对应网络class 需要在加载之前已经定义
the_model = torch.load(PATH)

#使用GPU时
model = TheModelClass().cuda()
model = nn.DataParallel(model).cuda()
model.load_state_dict(torch.load(best_model_path))
model.eval()

# 预测图像的类别
output = model(input)

4.网络构建

import torch

import torch.nn as nn

class SimpleNet(nn.Module):

    def __init__(self):
        super(SimpleNet, self).__init__()
        # 四个卷积层用于提取特征
        # 3 input channel image 160x160, 8 output channel image 80x80
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=8, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # 3 input channel image 80x80, 16 output channel image 40x40
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # 8 input channel image 40x40, 32 output channel image 20x20
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # 16 input channel image 20x20, 64 output channel image 10x10
        self.conv3 = nn.Sequential(
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # 分类
        self.classifier = nn.Sequential(
            nn.Linear(64 * 10 * 10, 6)
        )

    def forward(self, x):
        x = self.conv(x)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = x.view(-1, 64 * 10 * 10)
        x = self.classifier(x)
        return x

5.训练过程

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler

model = SimpleNet()
#是否使用gpu
if use_gpu:
        model = torch.nn.DataParallel(model)
        model.to(torch.device('cuda'))
    else:
        model.to(torch.device('cpu'))
确定loss Function
criterion = nn.CrossEntropyLoss()

确定优化方法
optimizer_ft = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=0.00004)
optimizer_ft = optim.Adam(model.parameters(), lr=args.lr)

学习率衰减
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=1, gamma=0.98)

device = torch.device('cuda' if use_gpu else 'cpu')

#test，log之类的没写，可以自己加，复制是不能直接跑的，只是介绍流程
for epoch in need_epoch:
    exp_lr_scheduler (epoch)
    model.train() #训练模式，batchNormal和dropout会跑，model.eval()时不跑
    # 在多个batch上依次处理数据(Iterate over data)
    for i, (inputs, labels) in enumerate(dataloders[phase]):
        inputs = inputs.to(device)
        labels = labels.to(device)

        # 梯度置零(zero the parameter gradients)
        optimizer.zero_grad()

         # 前向传播(forward)
         # 训练模式下才记录梯度以进行反向传播(track history if only in train)
         outputs = model(inputs)
         _, preds = torch.max(outputs, 1)
         loss = criterion(outputs, labels)
         # 训练模式下进行反向传播与梯度下降(backward + optimize only if in training phase)
         loss.backward()
         optimizer.step()

         # 统计损失和准确率(statistics)
         running_loss += loss.item() * inputs.size(0)
         running_corrects += torch.sum(preds == labels.data)

     batch_loss = running_loss / (i * args.batch_size + inputs.size(0))
     batch_acc = running_corrects.double() / (i * args.batch_size + inputs.size(0))

 epoch_loss = running_loss / dataset_sizes
 epoch_acc = running_corrects.double() / dataset_sizes

# 深拷贝模型(deep copy the model)
if phase == 'test' and epoch_acc > best_acc:
    best_acc = epoch_acc
    best_model_wts = copy.deepcopy(model.state_dict())

6.tensorboard 

安装 pip install tensorboard==2.0.0 (固定安装2.0版本，因为torch版本未到2.1，后面如果上去了可以自己改 )

# from tensorboardX import SummaryWriter
from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter(log_dir='log') #存储log的位置

writer.add_scalar(phase + '/Loss', epoch_loss, epoch) #数值型（表title，y轴数值，x轴数值）
writer.add_image(phase + '/Loss', image, epoch) #数值型（表title，图像值，x轴数值）

writer.close()

查看tensorboard：在存储log的位置运行命令行  tensorboard --logdir runs

 

 

注：有参考网上一个教程，具体链接找不到了。。。。