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)

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

 

 

注：有參考網上一個教程，具體鏈接找不到了。。。。