pytorch应用之——纸币识别（一）

最近参加了tinymind比赛，比赛内容是实现纸币识别，精度高于99.9%。
我用pytorch实现了这个要求，并记录在这里供大家跟自己参考回顾。

数据集https://pan.baidu.com/s/1oZ4yMlTU3YwnX9KR0huzyQ 提取码：k7bc

我分几个部分跟大家讲解：
1：算法思想（这个很简单）
2：代码详解

一：算法思想

这里数据集一共有39620张，而且背景单一，所以纸币面值的识别不是一个很难的问题。我用resnet18（自己稍微改了一些结构，影响不大）去训练这个数据集，迭代24次可以达到99.96%的精度。但这里要注意resnet不需要用预训练模型，然后resnet的最后的全连接改成number=9（因为数据集的label只有9种）。细节我会在代码中给大家介绍。

二.代码详解

代码部分我主要分以下几点讲解：
1：数据加载预处理
2：model修改
3：train过程。

1.数据加载预处理

class MyDataset(Dataset):
    def __init__(self, path):
        self.filenames = []
        self.labels = []
        tags = {"0.1": 0,
                "0.2": 1,
                "0.5": 2,
                "1": 3,
                "2": 4,
                "5": 5,
                "10": 6,
                "50": 7,
                "100": 8}
        with open(path, 'r') as file:
            for line in file:
                line = line.strip()
                words = line.split(",")
                filename = words[0].strip()
                filename = os.path.join('./data/train_data/', filename)
                
                self.filenames.append(filename)
                self.labels.append(label)

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

    def __getitem__(self, index):
        img_path = self.filenames[index]
      #  label = np.zeros(9)
      #  y = self.labels[index]
        label=self.labels[index]
        img = cv2.imread(img_path)  # BGR
        img = cv2.resize(img, (224, 224), cv2.INTER_CUBIC)
        assert img is not None, 'File Not Found ' + img_path
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
        img = np.ascontiguousarray(img, dtype=np.float32)  # uint8 to float32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        return torch.from_numpy(img), label, img_path

首先，我们需要将数据集加载到模型里面。初始化中从我们准备好的csv里获取图片路径和标签，存入filenames里，filennames是一个list。这里我们主要看getitem函数。

index是图片的索引，将对应索引的图片地址赋值给img_path，label也是相同操作，然后通过cv2来读取目标地址的图片。

将图片都resize成224，如果没有经过resize程序会报错，然后transpose以及归一，把图片变成pytorch可以处理的形式。

最后，返回torch（img），label，以及路径。

当MyDataset构建好，剩下的操作就交给DataLoder,在DataLoder中，会触发MyDataset中的getitem函数读取一张图片的数据和标签，并拼接成一个batch返回，作为模型真正的输入。

2.model修改
model_ft = models.resnet18(pretrained=False)，我用的是resnet18，这里我没有使用预训练模型，因为我的数据集比较简单且样本足够多，所以就直接训练resnet。但resnet在pytorch中默认的是训练imagenet数据集，它的输出是1000个类，我们需要把最后的输出改成9。

num_ftrs = model_ft.fc.in_features，把resnet18最后的fc层的in_features(输入节点数)赋值给num_ftrs。

model_ft.fc = nn.Linear(num_ftrs, 9)，将全连接的输出节点数改为9。这样我们就完成了model的修改。

criterion = nn.CrossEntropyLoss()，图像分类当然使用交叉熵损失函数。

optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)使用随机梯度下降的优化方式来BP。

exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)规划学习率，step_size=7表示每7个epochs调整学习率，衰减因子为0.1。

3.train

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()
    best_acc = 0.0
    weights = 'weights'
    latest = weights
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train']:
            if phase == 'train':
                scheduler.step()
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
           # for inputs, labels in enumerate(dataloader):
            for i, (inputs, labels, _) in enumerate(dataloader):
                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)
                if i % 50 == 0:
                    torch.save(model.state_dict(), latest)
                    print('epoch: {}, loss: {:.4}'.format(i, loss.data.item()))

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

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))

            # deep copy the model
            if  epoch_acc > best_acc:
                best_acc = epoch_acc


        print()

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))
    torch.save(model.state_dict(), latest)
    return model

我们将模型，损失函数，优化器，规划学习率，epochs传入函数。latest设置为保存模型参数的路径。

进入第一个epoch的循环，将模型设置为train，for i, (inputs, labels, _) in enumerate(dataloader):我们从dataloader这个list中把i索引，inputs图像数据，labels标签取出。把inputs放入模型，得到out，再与labels求loss。最后通过loss.backward()进行bp。

torch.save(model.state_dict(), latest)保存模型参数到指定路径。