train.py

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import os
import random
import argparse
from torchvision import transforms
from torch.autograd import Variable
import cv2
import numpy as np


class BaseLoader(torch.utils.data.Dataset):
    def __init__(self, triplets, transform=None):
        self.triplets = triplets
        self.transform = transform

    def __getitem__(self, index):
        img1_pth, img2_pth, img3_pth = self.triplets[index]
        img1 = cv2.imread(img1_pth)
        img2 = cv2.imread(img2_pth)
        img3 = cv2.imread(img3_pth)

        try:
            img1 = cv2.resize(img1, (args.picture_resize, args.picture_resize))
        except Exception as e:
            img1 = np.zeros((args.picture_resize, args.picture_resize, 3), dtype=np.uint8)

        try:
            img2 = cv2.resize(img2, (args.picture_resize, args.picture_resize))
        except Exception as e:
            img2 = np.zeros((args.picture_resize, args.picture_resize, 3), dtype=np.uint8)

        try:
            img3 = cv2.resize(img3, (args.picture_resize, args.picture_resize))
        except Exception as e:
            img3 = np.zeros((args.picture_resize, args.picture_resize, 3), dtype=np.uint8)

        if self.transform is not None:
            img1 = self.transform(img1)
            img2 = self.transform(img2)
            img3 = self.transform(img3)

        return img1, img2, img3

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


class BaseCnn(nn.Module):
    def __init__(self):
        super(BaseCnn, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, 3),
            nn.MaxPool2d(2)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 128, 3),
            nn.MaxPool2d(2)
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(128, 128, 3),
            nn.MaxPool2d(2)
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = x.view(x.size(0), -1)
        x = F.normalize(x, p=2, dim=1)
        return x


class SiameseNet(nn.Module):
    def __init__(self):
        super(SiameseNet, self).__init__()
        self.base = BaseCnn()

    def forward(self, x1, x2, x3):
        x1 = self.base(x1)
        x2 = self.base(x2)
        x3 = self.base(x3)

        return x1, x2, x3


class BaseDset(object):
    def __init__(self):
        self.__base_path = ""

        self.__train_set = {}
        self.__test_set = {}
        self.__train_keys = []
        self.__test_keys = []

    def load(self, base_path):
        """加载训练和测试数据集，将类别和路径存储"""
        self.__base_path = base_path
        train_dir = os.path.join(self.__base_path, 'train')
        test_dir = os.path.join(self.__base_path, 'test')

        self.__train_set = {}
        self.__test_set = {}
        self.__train_keys = []
        self.__test_keys = []

        for class_id in os.listdir(train_dir):
            class_dir = os.path.join(train_dir, class_id)
            self.__train_set[class_id] = []
            self.__train_keys.append(class_id)

            for img_name in os.listdir(class_dir):
                img_path = os.path.join(class_dir, img_name)
                self.__train_set[class_id].append(img_path)

        for class_id in os.listdir(test_dir):
            class_dir = os.path.join(test_dir, class_id)
            self.__test_set[class_id] = []
            self.__test_keys.append(class_id)
            for img_name in os.listdir(class_dir):
                img_path = os.path.join(class_dir, img_name)
                self.__test_set[class_id].append(img_path)

        return len(self.__train_keys), len(self.__test_keys)

    # 获取三元组
    def getTriplet(self, split='train'):
        # 默认选取训练集
        if split == 'train':
            dataset = self.__train_set
            keys = self.__train_keys
        else:
            dataset = self.__test_set
            keys = self.__test_keys

        # 随机指定两个正负类别，确保二者不一致
        pos_idx = random.randint(0, len(keys) - 1)
        while True:
            neg_idx = random.randint(0, len(keys) - 1)
            if pos_idx != neg_idx:
                break
        # 选定一个原始样本
        pos_anchor_img_idx = random.randint(0, len(dataset[keys[pos_idx]]) - 1)
        # 随机选择一个正样本，保证二者不一致
        while True:
            pos_img_idx = random.randint(0, len(dataset[keys[pos_idx]]) - 1)
            if pos_anchor_img_idx != pos_img_idx:
                break
        # 随机选择一个负样本
        neg_img_idx = random.randint(0, len(dataset[keys[neg_idx]]) - 1)

        # 生成三元组
        pos_anchor_img = dataset[keys[pos_idx]][pos_anchor_img_idx]
        pos_img = dataset[keys[pos_idx]][pos_img_idx]
        neg_img = dataset[keys[neg_idx]][neg_img_idx]

        return pos_anchor_img, pos_img, neg_img


def train(data, model, criterion, optimizer, epoch):
    print("******** Training ********")
    total_loss = 0
    model.train()
    for batch_idx, img_triplet in enumerate(data):
        # 提取三元组数据
        anchor_img, pos_img, neg_img = img_triplet
        anchor_img, pos_img, neg_img = anchor_img.to(device), pos_img.to(device), neg_img.to(device)
        anchor_img, pos_img, neg_img = Variable(anchor_img), Variable(pos_img), Variable(neg_img)
        # 分别获得三个编码，表示原始样本、正样本、负样本
        E1, E2, E3 = model(anchor_img, pos_img, neg_img)
        # 计算二者之间的欧式距离
        dist_E1_E2 = F.pairwise_distance(E1, E2, 2)
        dist_E1_E3 = F.pairwise_distance(E1, E3, 2)

        target = torch.FloatTensor(dist_E1_E2.size()).fill_(-1)
        target = target.to(device)
        target = Variable(target)

        loss = criterion(dist_E1_E2, dist_E1_E3, target)
        total_loss += loss

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # 打印损失
        log_step = args.train_log_step
        if (batch_idx % log_step == 0) and (batch_idx != 0):
            print('Train Epoch: {} [{}/{}] \t Loss: {:.4f}'.format(epoch, batch_idx, len(data), total_loss / log_step))
            total_loss = 0
    print("****************")


def test(data, model, criterion):
    print("******** Testing ********")
    with torch.no_grad():
        model.eval()
        accuracies = [0, 0, 0]
        acc_threshes = [0, 0.2, 0.5]
        total_loss = 0
        for batch_idx, img_triplet in enumerate(data):
            anchor_img, pos_img, neg_img = img_triplet
            anchor_img, pos_img, neg_img = anchor_img.to(device), pos_img.to(device), neg_img.to(device)
            anchor_img, pos_img, neg_img = Variable(anchor_img), Variable(pos_img), Variable(neg_img)
            E1, E2, E3 = model(anchor_img, pos_img, neg_img)
            dist_E1_E2 = F.pairwise_distance(E1, E2, 2)
            dist_E1_E3 = F.pairwise_distance(E1, E3, 2)

            target = torch.FloatTensor(dist_E1_E2.size()).fill_(-1)
            target = target.to(device)
            target = Variable(target)

            loss = criterion(dist_E1_E2, dist_E1_E3, target)
            total_loss += loss

            for i in range(len(accuracies)):
                prediction = (dist_E1_E3 - dist_E1_E2 - args.margin * acc_threshes[i]).cpu().data
                prediction = prediction.view(prediction.numel())
                prediction = (prediction > 0).float()
                batch_acc = prediction.sum() * 1.0 / prediction.numel()
                accuracies[i] += batch_acc
        print('Test Loss: {}'.format(total_loss / len(data)))
        for i in range(len(accuracies)):
            # 0%等价于准确率，其余是更严格的指标
            print(
                'Test Accuracy with diff = {}% of margin: {:.4f}'.format(acc_threshes[i] * 100,
                                                                         accuracies[i] / len(data)))
    print("****************")

    return accuracies[-1]


def main():
    # random_seed
    torch.manual_seed(718)
    torch.cuda.manual_seed(718)

    data_path = r'./characters'
    # data_path = r'./characters'
    dset_obj = BaseDset()
    dset_obj.load(data_path)

    train_triplets = []
    test_triplets = []

    for i in range(args.num_train_samples):
        pos_anchor_img, pos_img, neg_img = dset_obj.getTriplet()
        train_triplets.append([pos_anchor_img, pos_img, neg_img])
    for i in range(args.num_test_samples):
        pos_anchor_img, pos_img, neg_img = dset_obj.getTriplet(split='test')
        test_triplets.append([pos_anchor_img, pos_img, neg_img])
    loader = BaseLoader
    model = SiameseNet()
    model.to(device)

    criterion = torch.nn.MarginRankingLoss(margin=args.margin)
    optimizer = optim.Adam(model.parameters(), lr=args.lr)

    best_acc_of_50_margin = 0
    best_epoch = 0

    for epoch in range(1, args.epochs + 1):
        # 初始化数据加载器
        # 加载三元组
        train_data_loader = torch.utils.data.DataLoader(
            loader(train_triplets,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize(0, 1)
                   ])),
            batch_size=args.batch_size, shuffle=True)
        test_data_loader = torch.utils.data.DataLoader(
            loader(test_triplets,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize(0, 1)
                   ])),
            batch_size=args.batch_size, shuffle=True)
        train(train_data_loader, model, criterion, optimizer, epoch)
        acc_of_50_margin = test(test_data_loader, model, criterion)

        model_to_save = {
            "epoch": epoch + 1,
            'state_dict': model.state_dict(),
        }

        if acc_of_50_margin > best_acc_of_50_margin:
            best_acc_of_50_margin = acc_of_50_margin
            best_epoch = epoch

            if not args.disable_save_best_ckp:
                result_path = os.path.join(args.result_dir)
                file_name = os.path.join(args.result_dir, "best_checkpoint" + ".pt")
                if not os.path.exists(result_path):
                    os.makedirs(result_path)
                save_checkpoint(model_to_save, file_name)

        if (epoch % args.ckp_freq == 0) and not args.disable_save_ckp:
            result_path = os.path.join(args.result_dir)
            file_name = os.path.join(args.result_dir, "checkpoint_" + str(epoch) + ".pt")
            if not os.path.exists(result_path):
                os.makedirs(result_path)
            save_checkpoint(model_to_save, file_name)
    print("Training is done.")
    print(f"The best epoch of acc50, which is {best_acc_of_50_margin * 100}%, is {best_epoch}.")


def save_checkpoint(state, file_name):
    torch.save(state, file_name)


if __name__ == '__main__':
    # 超参数
    parser = argparse.ArgumentParser(description='PyTorch Siamese Example')
    parser.add_argument('--result_dir', default='output', type=str,
                        help='Directory to store results')
    parser.add_argument('--epochs', type=int, default=10, metavar='N',
                        help='number of epochs to train (default: 10)')

    parser.add_argument("--disable_save_ckp", default=False, action='store_true',
                        help="disable to save checkpoint frequently")
    parser.add_argument('--ckp_freq', type=int, default=5, metavar='N',
                        help='Checkpoint Frequency (default: 1)')
    parser.add_argument("--disable_save_best_ckp", default=False, action='store_true',
                        help="disable to save best checkpoint")

    parser.add_argument('--train_log_step', type=int, default=500, metavar='M',
                        help='Number of iterations after which to log the loss')
    parser.add_argument('--margin', type=float, default=1.0, metavar='M',
                        help='margin for triplet loss (default: 1.0)')

    parser.add_argument('--batch_size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--lr', type=float, default=0.0001, metavar='LR',
                        help='learning rate (default: 0.0001)')
    parser.add_argument('--dataset', type=str, default='mnist', metavar='M',
                        help='Dataset (default: mnist)')
    parser.add_argument('--picture_resize', type=int, default=200, metavar='M',
                        help='size of the picture to reset (default: 200)')
    parser.add_argument('--num_train_samples', type=int, default=50000, metavar='M',
                        help='number of training samples (default: 50000)')
    parser.add_argument('--num_test_samples', type=int, default=10000, metavar='M',
                        help='number of test samples (default: 10000)')

    global args, device
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    args = parser.parse_args()

    main()