0040-pytorch入门-猫狗二分类检测模型

淡淡的烟草味﹌ 2023-01-12 14:59 92阅读 0赞

### ImageDataset.py ###

# -*- coding: utf-8 -*-
    # @Time    : 2021/1/28 17:23
    # @Author  : Johnson
    '''
    负责训练以及测试数据的读取
    '''
    from torchvision import transforms,datasets
    import os
    import torch
    from PIL import Image
    
    def readImg(path):
        '''
        用于替代ImageFolder的默认读取图片函数，以读取单通道图片
        :param path:
        :return:
        '''
        return Image.open(path)
    
    def ImageDataset(args):
        #数据增强以及归一化
        #图片都是190*100，训练时随机裁剪90*90，测试时裁剪中间的90*90
        data_transforms = {
            'train': transforms.Compose([
                transforms.RandomCrop(90),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize([0.5], [0.5])
            ]),
            'test': transforms.Compose([
                transforms.CenterCrop(90),
                transforms.ToTensor(),
                transforms.Normalize([0.5], [0.5])
            ]),
        }
        data_dir = args.data_dir
        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

# simpleNet.py #

# -*- coding: utf-8 -*-
    # @Time    : 2021/1/28 17:35
    # @Author  : Johnson
    '''
    简单的用于分类的网络
    '''
    
    import torch
    
    import torch.nn as nn
    
    class SimpleNet(nn.Module):
    
        def __init__(self):
            super(SimpleNet, self).__init__()
            # 三个卷积层用于提取特征
            # 1 input channel image 90x90, 8 output channel image 44x44
            self.conv1 = nn.Sequential(
                nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3, stride=1, padding=0),
                nn.ReLU(),
                nn.MaxPool2d(2)
            )
            # 8 input channel image 44x44, 16 output channel image 22x22
            self.conv2 = nn.Sequential(
                nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride=1, padding=1),
                nn.ReLU(),
                nn.MaxPool2d(2)
            )
            # 16 input channel image 22x22, 32 output channel image 10x10
            self.conv3 = nn.Sequential(
                nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=0),
                nn.ReLU(),
                nn.MaxPool2d(2)
            )
            # 分类
            self.classifier = nn.Sequential(
                nn.Linear(32 * 10 * 10, 3)
            )
    
        def forward(self, x):
            x = self.conv1(x)
            x = self.conv2(x)
            x = self.conv3(x)
            x = x.view(-1, 32 * 10 * 10)
            x = self.classifier(x)
            return x

### train.py ###

# -*- coding: utf-8 -*-
    # @Time    : 2021/1/28 17:35
    # @Author  : Johnson
    '''
    进行训练
    '''
    
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.optim import lr_scheduler
    import time
    import os
    import json
    from math import ceil
    import argparse
    import copy
    from ImageDataset import ImageDataset
    from  simpleNet import SimpleNet
    from tensorboardX import SummaryWriter
    
    writer = SummaryWriter(log_dir='log')
    
    
    def train_model(args, model, criterion, optimizer, scheduler, num_epochs, dataset_sizes, use_gpu):
        since = time.time()
    
        best_model_wts = copy.deepcopy(model.state_dict())
        best_acc = 0.0
    
        device = torch.device('cuda' if use_gpu else 'cpu')
    
        for epoch in range(args.start_epoch, num_epochs):
            # 每一个epoch中都有一个训练和一个验证过程(Each epoch has a training and validation phase)
            for phase in ['train', 'test']:
                if phase == 'train':
                    scheduler.step(epoch)
                    # 设置为训练模式(Set model to training mode)
                    model.train()
                else:
                    # 设置为验证模式(Set model to evaluate mode)
                    model.eval()
    
                running_loss = 0.0
                running_corrects = 0
    
                tic_batch = time.time()
    
                # 在多个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)
                    with torch.set_grad_enabled(phase == 'train'):
                        outputs = model(inputs)
                        _, preds = torch.max(outputs, 1)
                        loss = criterion(outputs, labels)
                        # 训练模式下进行反向传播与梯度下降(backward + optimize only if in training phase)
                        if phase == 'train':
                            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))
    
                    if phase == 'train' and (i + 1) % args.print_freq == 0:
                        print(
                            '[Epoch {}/{}]-[batch:{}/{}] lr:{:.6f} {} Loss: {:.6f}  Acc: {:.4f}  Time: {:.4f} sec/batch'.format(
                                epoch + 1, num_epochs, i + 1, ceil(dataset_sizes[phase] / args.batch_size),
                                scheduler.get_lr()[0], phase, batch_loss, batch_acc,
                                (time.time() - tic_batch) / args.print_freq))
                        tic_batch = time.time()
    
                epoch_loss = running_loss / dataset_sizes[phase]
                epoch_acc = running_corrects.double() / dataset_sizes[phase]
    
                if epoch == 0:
                    os.remove('result.txt')
                with open('result.txt', 'a') as f:
                    f.write('Epoch:{}/{} {} Loss: {:.4f} Acc: {:.4f} \n'.format(epoch + 1, num_epochs, phase, epoch_loss,
                                                                                epoch_acc))
    
                print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
    
                writer.add_scalar(phase + '/Loss', epoch_loss, epoch)
                writer.add_scalar(phase + '/Acc', epoch_acc, epoch)
    
            if (epoch + 1) % args.save_epoch_freq == 0:
                if not os.path.exists(args.save_path):
                    os.makedirs(args.save_path)
                torch.save(model.state_dict(), os.path.join(args.save_path, "epoch_" + str(epoch) + ".pth"))
    
            # 深拷贝模型(deep copy the model)
            if phase == 'test' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
    
        # 将model保存为graph
        writer.add_graph(model, (inputs,))
    
        time_elapsed = time.time() - since
        print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
        print('Best val Accuracy: {:4f}'.format(best_acc))
    
        # 载入最佳模型参数(load best model weights)
        model.load_state_dict(best_model_wts)
        return model
    
    
    if __name__ == '__main__':
    
        parser = argparse.ArgumentParser(description='classification')
        # 图片数据的根目录(Root catalog of images)
        parser.add_argument('--data-dir', type=str, default='images')
        parser.add_argument('--batch-size', type=int, default=16)
        parser.add_argument('--num-epochs', type=int, default=150)
        parser.add_argument('--lr', type=float, default=0.045)
        parser.add_argument('--num-workers', type=int, default=4)
        parser.add_argument('--print-freq', type=int, default=1)
        parser.add_argument('--save-epoch-freq', type=int, default=1)
        parser.add_argument('--save-path', type=str, default='output')
        parser.add_argument('--resume', type=str, default='', help='For training from one checkpoint')
        parser.add_argument('--start-epoch', type=int, default=0, help='Corresponding to the epoch of resume')
        args = parser.parse_args()
    
        # read data
        dataloders, dataset_sizes, class_names = ImageDataset(args)
    
        with open('class_names.json', 'w') as f:
            json.dump(class_names, f)
    
        # use gpu or not
        use_gpu = torch.cuda.is_available()
        print("use_gpu:{}".format(use_gpu))
    
        # get model
        model = SimpleNet()
    
        if args.resume:
            if os.path.isfile(args.resume):
                print(("=> loading checkpoint '{}'".format(args.resume)))
                model.load_state_dict(torch.load(args.resume))
            else:
                print(("=> no checkpoint found at '{}'".format(args.resume)))
    
        if use_gpu:
            model = torch.nn.DataParallel(model)
            model.to(torch.device('cuda'))
        else:
            model.to(torch.device('cpu'))
    
        # 用交叉熵损失函数(define loss function)
        criterion = nn.CrossEntropyLoss()
    
        # 梯度下降(Observe that all parameters are being optimized)
        optimizer_ft = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=0.00004)
    
        # Decay LR by a factor of 0.98 every 1 epoch
        exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=1, gamma=0.98)
    
        model = train_model(args=args,
                            model=model,
                            criterion=criterion,
                            optimizer=optimizer_ft,
                            scheduler=exp_lr_scheduler,
                            num_epochs=args.num_epochs,
                            dataset_sizes=dataset_sizes,
                            use_gpu=use_gpu)
    
        torch.save(model.state_dict(), os.path.join(args.save_path, 'best_model_wts.pth'))
    
        writer.close()

### test.py ###

# -*- coding: utf-8 -*-
    # @Time    : 2021/1/28 17:36
    # @Author  : Johnson
    '''
    测试分类
    '''
    
    from PIL import Image
    from torchvision import transforms
    import torch
    from torch.autograd import Variable
    import os
    import json
    from simpleNet import SimpleNet
    
    def predict_image(model, image_path):
        image = Image.open(image_path)
    
        # 测试时截取中间的90x90
        transformation1 = transforms.Compose([
            transforms.CenterCrop(90),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5])
    
        ])
    
        # 预处理图像
        image_tensor = transformation1(image).float()
    
        # 额外添加一个批次维度，因为PyTorch将所有的图像当做批次
        image_tensor = image_tensor.unsqueeze_(0)
    
        if torch.cuda.is_available():
            image_tensor.cuda()
    
        # 将输入变为变量
        input = Variable(image_tensor)
    
        # 预测图像的类别
        output = model(input)
    
        index = output.data.numpy().argmax()
    
        return index
    
    if __name__ == '__main__':
    
        best_model_path = './output/epoch_462.pth'
        model = SimpleNet()
        model.load_state_dict(torch.load(best_model_path))
        model.eval()
    
        with open('class_names.json', 'r') as f:
            class_names = json.load(f)
    
        img_path = './images/test/bubble/066.jpg'
        predict_class = class_names[predict_image(model, img_path)]
        print(predict_class)

### 参考链接 ###

*  [用pytorch训练图像分类器模型导出ONNX测试][pytorch_ONNX]
 *  [用pytorch实现一个简单的图像分类器][pytorch]

[pytorch_ONNX]: https://blog.csdn.net/dcrmg/article/details/102807575
[pytorch]: https://blog.csdn.net/JR_Chan/article/details/95641758