pytorch 车型分类代码

pytorch 车型分类代码 （1）

• 可以跑通，代码简洁
• 学习自用
• 持续更新

1 数据

1.1 文件配置

数据需要按照下面的方式组织，在数据文件夹中需要有train和val文件夹，然后每个文件夹内中，每个类别放在一个文件夹即可，可以直接利用pytorch自带的datasets.ImageFolder函数加载数据，image有点像imagenet 分类配置。* 这种凡方式也有弊端，一般业务场景数据是通过一个image_list.txt 文件加载，其中包含图片地址和label，调整类别映射的数字不方便*

1.2 数据增强

- 使用torch自带的transform函数，过程如下
- 随机裁剪
- 水平翻转
- 归一化（像素值编程0到1之间，并且转成torch张量）
- 标准化（减去均值，除以方差）

```python
train_transforms = transforms.Compose(
        [transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)),  # 随机裁剪到256*256
         transforms.RandomRotation(degrees=15),  # 随机旋转
         transforms.RandomHorizontalFlip(),  # 随机水平翻转
         transforms.CenterCrop(size=224),  # 中心裁剪到224*224
         transforms.ToTensor(),  # 转化成张量,#归一化
         transforms.Normalize([0.485, 0.456, 0.406],
                              [0.229, 0.224, 0.225])  # 标准化
         ])
```

1.3数据加载

1.3.1 使用ImageFolder加载数据

 train_datasets = datasets.ImageFolder(TRAIN_DATASET_DIR, transform=train_transforms)
	 train_dataloader = torch.utils.data.DataLoader(train_datasets,
	                                                   batch_size=TRAIN_BATCH_SIZE,
	                                                   shuffle=True,
	                                                   **kwargs)

• 查看数据类别映射
  • 通过 datasets.ImageFolder 可以查看类别映射 -

   print(train_datasets.class_to_idx)
    >> {'SUV': 0, 'bus': 1, 'family sedan': 2, 'fire engine': 3, 'heavy truck': 4, 'jeep': 5, 'minibus': 6, 'racing car': 7, 'taxi': 8, 'truck': 9}
   print(train_datasets.classes)
   >>['SUV', 'bus', 'family sedan', 'fire engine', 'heavy truck', 'jeep', 'minibus', 'racing car', 'taxi', 'truck']

   

1.3.2 自己定义数据结构

class MY_DATASET(Dataset):
    def __init__(self, train, train_data_list_txt, val_data_list_txt, transform,
                 target_transform):
        self.train = train
        self.train_data_list_txt = train_data_list_txt
        self.val_data_list_txt = val_data_list_txt
        self.transform = transform
        self.target_transform = target_transform
        if self.train:
            self.data_file = train_data_list_txt
        else:
            self.data_file = val_data_list_txt
        self.data_list = self.read_data_list(self.data_file)

    def read_data_list(self, data_path):
        print("data_path", data_path)
        assert os.path.exists(data_path)
        with open(data_path, "r") as fp:
            print("data_path", data_path)
            data_list = fp.readlines()
        return data_list

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        img_path, target = self.data_list[index].strip().split("  ")
        target = int(target)
        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        if not os.path.exists(img_path):
            print("{} is not exists!!!!!!!!!!!".format(img_path))
            raise
        img = cv2.imread(img_path)
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            self.target = self.target_transform(target)
        return img, target

    def __len__(self):
        return len(self.data_list)

1.4 数据结构

2 代码

"""
仅用于学习

********类别说明**********>
0，巴士，bus
1，出租车，taxi
2，货车，truck
3，家用轿车，family sedan
4，面包车，minibus
5，吉普车，jeep
6，运动型多功能车，SUV
7，重型货车，heavy truck
8，赛车，racing car
9，消防车，fire engine
"""
import torch
from torchvision import transforms
from torchvision import datasets
import matplotlib.pyplot as plt
import torch.nn.functional as F
import numpy as np
from  torchvision import models
import torch.nn as nn
from torch import optim

#****************setting*******************
NUM_CLASSES = 10
BATCH_SIZE = 32
NUM_EPOCHS= 25
#下载地址：https://download.pytorch.org/models/resnet50-19c8e357.pth
PRETRAINED_MODEL = './resnet50-19c8e357.pth'
MODEL_SAVE_PATH = 'trained_models/vehicle-10_record.pth'#数据集的存放位置
TRAIN_DATASET_DIR = r'F:/car_class10_data/train'
VALID_DATASET_DIR = r'F:/car_class10_data/val'
TRAIN_BATCH_SIZE = 128
TEST_BATCH_SIZE = 128
DROPOUT_RATE = 0.3
show_interval_num = 10
epochs = 20

#此处数据是分别放在10个文件夹中

#****************设置数据增强方式**************************
#针对训练集train_data
train_transforms = transforms.Compose(
        [transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)),#随机裁剪到256*256
        transforms.RandomRotation(degrees=15),#随机旋转
        transforms.RandomHorizontalFlip(),#随机水平翻转
        transforms.CenterCrop(size=224),#中心裁剪到224*224
        transforms.ToTensor(),#转化成张量,#归一化
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])#标准化
])
#针对测试集,test data,测试就不需要随机中心裁剪了，直接resize到224*224
test_valid_transforms = transforms.Compose(
        [transforms.Resize(224),
         transforms.ToTensor(),
         transforms.Normalize([0.485, 0.456, 0.406],
         [0.229, 0.224, 0.225])])
#*************************通过pytorch自带的dataload加载数据**************************
#关于dataload 可以查看 https://blog.csdn.net/weixin_40123108/article/details/85099449
# ImageFolder假设所有的文件按文件夹保存，每个文件夹下存储同一个类别的图，详细的可以去了解这个类ImageFolder，主要关注__getitem__函数，该函数会根据索引返回每张图和label
# 关于torch.utils.data.DataLoader,数据加载器，结合了数据集和取样器，并且可以提供多个线程处理数据集。
# 在训练模型时使用到此函数，用来把训练数据分成多个小组，此函数每次抛出一组数据。直至把所有的数据都抛出。就是做一个数据的初始化。
# 可以参考学习 https://zhuanlan.zhihu.com/p/28200166    https://www.jb51.net/article/184042.htm
train_datasets = datasets.ImageFolder(TRAIN_DATASET_DIR, transform=train_transforms)
train_dataloader = torch.utils.data.DataLoader(train_datasets, batch_size=TRAIN_BATCH_SIZE, shuffle=True)
train_data_size = len(train_datasets)
valid_datasets = datasets.ImageFolder(VALID_DATASET_DIR,transform=test_valid_transforms)
valid_dataloader = torch.utils.data.DataLoader(valid_datasets, batch_size=TEST_BATCH_SIZE, shuffle=True)
valid_data_size = len(valid_datasets)
#****************************可以通过运行test_data函数查看数据类型**************************************

def test_data():
    print("train_dataloade len", len(train_dataloader))
    for images, labels in train_dataloader:
        print(labels)
        print("label len", len(labels))
        img = images[0]
        img = img.numpy()
        img = np.transpose(img, (1, 2, 0))
        plt.imshow(img)
        plt.show()
        break
#*******************使用预训练模型 resnet50进行fineturn**************************
#修改最后一层fc
def resnet50():
    model = models.resnet50(pretrained=True)
    for param in model.parameters():
        param.requires_grad = False
    fc_inputs = model.fc.in_features
    model.fc = nn.Sequential(
        nn.Linear(fc_inputs, 256),
        nn.ReLU(),
        nn.Dropout(0.4),
        nn.Linear(256, 10),
        nn.LogSoftmax(dim=1)
    )
    return model
#********************定义损失函数和优化器*********************
loss_func = nn.NLLLoss()
optimizer = optim.Adam(resnet50.parameters())
#**********************定义训练和验证过程***************************
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % show_interval_num == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
#****************************定义main*****************************
def main():
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
    train_transforms = transforms.Compose(
        [transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)),  # 随机裁剪到256*256
         transforms.RandomRotation(degrees=15),  # 随机旋转
         transforms.RandomHorizontalFlip(),  # 随机水平翻转
         transforms.CenterCrop(size=224),  # 中心裁剪到224*224
         transforms.ToTensor(),  # 转化成张量,#归一化
         transforms.Normalize([0.485, 0.456, 0.406],
                              [0.229, 0.224, 0.225])  # 标准化
         ])
    # 针对测试集,test data,测试就不需要随机中心裁剪了，直接resize到224*224
    test_valid_transforms = transforms.Compose(
        [transforms.Resize(224),
         transforms.ToTensor(),
         transforms.Normalize([0.485, 0.456, 0.406],
                              [0.229, 0.224, 0.225])])
    train_datasets = datasets.ImageFolder(TRAIN_DATASET_DIR, transform=train_transforms)
    train_dataloader = torch.utils.data.DataLoader(train_datasets,
                                                   batch_size=TRAIN_BATCH_SIZE,
                                                   shuffle=True,
                                                   **kwargs)
    valid_datasets = datasets.ImageFolder(VALID_DATASET_DIR, transform=test_valid_transforms)
    valid_dataloader = torch.utils.data.DataLoader(valid_datasets,
                                                   batch_size=TEST_BATCH_SIZE,
                                                   shuffle=True,
                                                   **kwargs)
    model = resnet50().to(device)
    optimizer = optim.Adam(resnet50.parameters(), lr=0.001)
    #***********************print flops and params**************************
    for epoch in range(1, epochs + 1):
        train(model, device, train_dataloader, optimizer, epoch)
        test_acc = test(model, device, valid_dataloader)
        # report intermediate result
        print('test accuracy %g', test_acc)
    # report final result
    print('Final result is %g', test_acc)
def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            # sum up batch loss
            test_loss += F.nll_loss(output, target, reduction='sum').item()
            # get the index of the max log-probability
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    accuracy = 100. * correct / len(test_loader.dataset)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset), accuracy))
    return accuracy


if __name__ == "__main__":
    test_data()
    pass

3结果