Ptorch 搭建分類模型

1. 神經網絡模型

腳本可以選擇的網絡有：Mobilenet_v2, Googlenet, Inception_v3, resnet50, Densnet121。 當然也可以添加你自己的網絡。

 

2.模型訓練

a. 數據路徑下包含train和val兩個文件夾，文件夾下面存放所有類別的數據，一個類別一個文件夾。

b. resize是圖片壓縮尺寸，crop-size是圖片中心剪裁后輸入網絡的尺寸。

c. 如果要加載訓練過的模型，開啟pre，并設置model-path模型路徑。

d. 如果要使用Focal-loss，需要新建Focal-loss.py，并拷貝代碼，代碼我會在後面列出。

訓練代碼：

import torch
import torchvision
import torch.nn as nn
from torchvision import transforms
import copy
import FocalLoss
import argparse, os
import time, datetime

# import tqdm

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--net-name', dest='net_name', type=str, default='resnet50')
    parser.add_argument('--data', dest='data', type=str, default='./data/class', help='The data path!')
    parser.add_argument('--resize', dest='resize', type=int, default=300, help='img resize')
    parser.add_argument('--crop-size', dest='crop_size', type=int, default=224, help='crop resized img enter net!')
    parser.add_argument('--batch-size', dest='batch_size', type=int, default=16)
    parser.add_argument('--epochs', dest='epochs', type=int, default=100)
    parser.add_argument('--classes', dest='classes', type=int, default=3, help='class number')
    parser.add_argument('--save-path', dest='save_path', type=str, default='./model', help='save model path!')
    parser.add_argument('--pre', dest='pre_training', action='store_true', help='pre_training or not!')
    parser.add_argument('--model-path', dest='model_path', type=str, default='./model/epoch0_acc0.3676_loss1.0442.pt',
                        help='pre_training model path!')
    parser.add_argument('--focal-loss', dest='focal_loss', action='store_true', help='use focal loss!')
    parser.add_argument('--fe', dest='feature_extract', default=True,
                        help='Flag for feature extractiing, When False, wei finetune the whole mode, When True, we only update the reshaped layer paras!')
    args = parser.parse_args()
    return args


args = parse_args()

transform = {
    'train': transforms.Compose([transforms.Resize(args.resize),
                                 transforms.CenterCrop(args.crop_size),
                                 transforms.ToTensor(),
                                 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
    'val': transforms.Compose([transforms.Resize(args.resize),
                               transforms.CenterCrop(args.crop_size),
                               transforms.ToTensor(),
                               transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
}


def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            ## 設置 requires_grad=False 凍結參數,以便在backward()中不計算梯度.
            param.requires_grad = False


def Net(feature_extract, net_name):
    net = None
    if net_name in ['Mobilenet_v2', 'mobilenet_v2']:
        net = torchvision.models.mobilenet_v2(pretrained=True)
        set_parameter_requires_grad(net, feature_extract)
        classifier = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(net.last_channel, args.classes),
        )
        net.classifier = classifier

    elif net_name in ['googlenet', 'Googlenet']:
        net = torchvision.models.googlenet(pretrained=True)
        set_parameter_requires_grad(net, feature_extract)
        in_features = net.fc.in_features
        fc = nn.Linear(in_features, args.classes, bias=True)
        net.fc = fc

    elif net_name in ['inception_v3', 'Inception_v3']:
        net = torchvision.models.inception_v3(pretrained=True)
        set_parameter_requires_grad(net, feature_extract)
        num_ftrs = net.AuxLogits.fc.in_features
        AuxLogits = nn.Linear(num_ftrs, args.classes)
        net.AuxLogits.fc = AuxLogits
        num_ftrs = net.fc.in_features
        fc = nn.Linear(num_ftrs, args.classes)
        net.fc = fc

    elif net_name in ['resnet50', 'Resnet50']:
        net = torchvision.models.resnet50(pretrained=True)
        set_parameter_requires_grad(net, feature_extract)
        in_features = net.fc.in_features
        fc = nn.Linear(in_features, args.classes)
        net.fc = fc

    elif net_name in ['densenet121', 'Densnet121']:
        net = torchvision.models.densenet121(pretrained=True)
        set_parameter_requires_grad(net, feature_extract)
        in_features = net.classifier.in_features
        classifier = nn.Linear(in_features, args.classes)
        net.classifier = classifier

    else:
        assert net, 'please add yourself net or input right net name!'

    print(list(net.children()))
    net = net.to(device)
    return net


def train(net, data_loader, optim, criterion, exp_lr_scheduler, epochs, net_name):
    best_acc = 0
    best_model_wts = copy.deepcopy(net.state_dict())
    for epoch in range(epochs):
        print('Epoch{}/{}'.format(epoch, epochs - 1))

        for phase in ['train', 'val']:
            if phase == 'train':
                exp_lr_scheduler.step()
                net.train()
            if phase == 'val':
                net.eval()

            running_loss = 0
            running_corects = 0
            step = 1
            steps = len(data_loader[phase])
            for inputs, labels in data_loader[phase]:
                strat_time = time.time()
                inputs = inputs.to(device)
                labels = labels.to(device)

                optim.zero_grad()

                with torch.set_grad_enabled(phase == 'train'):

                    if net_name in ['inception_v3', 'Inception_v3'] and phase == 'train':
                        outputs, aux_outputs = net(inputs)
                        loss1 = criterion(outputs, labels)
                        loss2 = criterion(aux_outputs, labels)
                        loss = loss1 + 0.4 * loss2
                    else:
                        outputs = net(inputs)
                        loss = criterion(outputs, labels)
                    _, predict = torch.max(outputs, 1)

                    if phase == 'train':
                        loss.backward()
                        optim.step()

                    end_time = time.time()
                    residual_time = str(datetime.timedelta(seconds=(steps - step) * (end_time - strat_time)))[:-7]
                    print("\r%d/%d [%s>%s] -ETA: %s - loss: %4f\n" % (
                        step, steps, '=' * int(29 * step / steps), '.' * (29 - int(29 * step / steps)), residual_time, loss), end='', flush=True)
                    step += 1

                running_loss += loss.item() * inputs.size(0)
                running_corects += torch.sum(predict == labels.data)

            epoch_loss = running_loss / data_size[phase]
            epoch_acc = running_corects.double() / data_size[phase]

            print('{} Loss:{:.4f} acc:{:.4f}'.format(phase, epoch_loss, epoch_acc))
            if epoch_acc > best_acc and phase == 'val':
                best_acc = epoch_acc
                torch.save(net.state_dict(), os.path.join(args.save_path,
                                                          'epoch{}_acc{:.4f}_loss{:.4f}.pt'.format(epoch, epoch_acc, epoch_loss)))
                best_model_wts = copy.deepcopy(net.state_dict())

    print('Best val acc', best_acc)
    net.load_state_dict(best_model_wts)
    torch.save(net.state_dict(), os.path.join(args.save_path, 'best_model_acc{:.4f}.pt'.format(best_acc)))


if __name__ == '__main__':
    # load data
    imgs_datasets = {x: torchvision.datasets.ImageFolder(os.path.join(args.data, x), transform=transform[x]) for x in ['train', 'val']}
    data_loader = {
        x: torch.utils.data.DataLoader(imgs_datasets[x], batch_size=args.batch_size, shuffle=True if x == 'train' else False, num_workers=0) for x in ['train', 'val']}
    data_size = {x: len(imgs_datasets[x]) for x in ['train', 'val']}
    img_class = imgs_datasets['train'].classes

    net = Net(args.feature_extract, args.net_name)

    # load model
    if args.pre_training:
        print('load model:{}'.format(args.model_path))
        net.load_state_dict(torch.load(args.model_path))

    # Observe that all parameters are being optimized
    if args.feature_extract:
        params_to_update = []
        for name, param in net.named_parameters():
            if param.requires_grad == True:
                params_to_update.append(param)
    else:
        params_to_update = net.parameters()
    optim = torch.optim.Adam(params=params_to_update, lr=0.001)
    # optim = torch.optim.SGD(params=net.parameters(), lr=0.001, momentum=0.9)

    # Loss function
    criterion = FocalLoss() if args.focal_loss else nn.CrossEntropyLoss()
    print('criterion: ', criterion)

    # Decay LR by a factor of 0.1 every 7 epochs
    exp_lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer=optim, step_size=7, gamma=0.1)
    train(net, data_loader, optim, criterion, exp_lr_scheduler, args.epochs, args.net_name)

Focal-loss代碼：

import torch
import torch.nn as nn
import torch.nn.functional as F


class FocalLoss(nn.Module):  # 1d and 2d

    def __init__(self, gamma=2, size_average=True):
        super(FocalLoss, self).__init__()
        self.gamma = gamma
        self.size_average = size_average

    def forward(self, logit, target, class_weight=None, type='softmax'):
        target = target.view(-1, 1).long()
        if type == 'sigmoid':
            if class_weight is None:
                class_weight = [1] * 2  # [0.5, 0.5]

            prob = torch.sigmoid(logit)
            prob = prob.view(-1, 1)
            prob = torch.cat((1 - prob, prob), 1)
            select = torch.FloatTensor(len(prob), 2).zero_().cpu()
            select.scatter_(1, target, 1.)

        elif type == 'softmax':
            B, C = logit.size()
            if class_weight is None:
                class_weight = [1] * C  # [1/C]*C

            # logit   = logit.permute(0, 2, 3, 1).contiguous().view(-1, C)
            prob = F.softmax(logit, 1)
            select = torch.FloatTensor(len(prob), C).zero_().cpu()
            select.scatter_(1, target, 1.)

        class_weight = torch.FloatTensor(class_weight).cpu().view(-1, 1)
        class_weight = torch.gather(class_weight, 0, target)

        prob = (prob * select).sum(1).view(-1, 1)
        prob = torch.clamp(prob, 1e-8, 1 - 1e-8)
        batch_loss = - class_weight * (torch.pow((1 - prob), self.gamma)) * prob.log()

        if self.size_average:
            loss = batch_loss.mean()
        else:
            loss = batch_loss

        return loss

3. 模型測試

import torch
import torchvision
import torch.nn as nn
from torchvision import transforms
from collections import Counter

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')


def Load_Data(img_size, imgs_path):
    transform = transforms.Compose([transforms.Resize(img_size),
                                    transforms.ToTensor(),
                                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])

    testset = torchvision.datasets.ImageFolder(root=imgs_path, transform=transform)
    testloader = torch.utils.data.DataLoader(testset, batch_size=8, shuffle=False, num_workers=0)
    return (testset, testloader)


def Load_Model(model_path):
    ## 訓練時只保存權重參數，需加載網絡
    # net = torchvision.models.mobilenet_v2(pretrained=False)
    # classifier = nn.Sequential(
    #     nn.Dropout(0.2),
    #     nn.Linear(net.last_channel, 3),
    # )
    # net.classifier = classifier
    # model = net.to(device)
    # state_dict = torch.load(model_path)
    # model.load_state_dict(state_dict)

    # 訓練時保存了整個模型
    model = torch.load(model_path)
    return model


def test(data, model):
    testset, testloader = data
    labels_list = testset.targets
    classes_list = testset.classes
    dict_right = {}
    dict_count = {}
    for i in range(len(classes_list)):
        dict_right[i] = 0
        dict_count[i] = Counter(labels_list)[i]
    with torch.no_grad():
        for data in testloader:
            inputs, labels = data
            inputs.to(device)
            labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            print('predicted:', predicted, 'labels:', labels)
            for i in range(len(labels)):
                if predicted[i] == labels[i]:
                    dict_right[int(labels[i])] += 1

    for i in range(len(classes_list)):
        acc = dict_right[i] / dict_count[i]
        print('Class {} Acc: {:.4f}'.format(classes_list[i], acc))


if __name__ == '__main__':
    model_path = r'./model/epoch0_acc0.9044_loss0.2180.pt'    # 模型路徑
    imgs_path = './data/class/val'                            # 測試集路徑
    img_size = 224                                            # 輸入網絡尺寸
    model = Load_Model(model_path)
    testdata = Load_Data(img_size=img_size, imgs_path=imgs_path)
    test(testdata, model)

4. 單張圖片預測

import os
import torch
from torchvision import transforms
from PIL import Image
import time

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
img_size = 224
model = torch.load('./model/epoch0_acc0.9044_loss0.2180.pt')

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


def predict_image(img_path):
    start_time = time.time()
    img = Image.open(img_path)
    input = data_transforms(img).unsqueeze(0)
    input.to(device)
    output = model(input)
    _, predict = torch.max(output, 1)
    end_time = time.time()
    print('use_time', end_time - start_time)
    print('precicted classes: ', predict.numpy()[0])


while True:
    img_path = input('please input img_path:')
    if not os.path.exists(img_path) and img_path != 'q':
        print("The path error, Try again!")
        continue
    if img_path == 'q':
        break
    predict_image(img_path)

參考文章：https://blog.csdn.net/weixin_40123108/article/details/85714030