def main():
# 训练集图片读取
class TrainDataset(Dataset):
    def __init__(self, label_list, transform=None, target_transform=None, loader=default_loader):
        imgs = []
        for index, row in label_list.iterrows():
            imgs.append((row['img_path'], row['labels']))
        self.imgs = imgs
        self.transform = transform
        self.target_transform = target_transform
        self.loader = loader

    def __getitem__(self, index):
        filename, label = self.imgs[index]
        img = self.loader(filename)
        if self.transform is not None:
            img = self.transform(img)
        return img, label

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

# 验证集图片读取
class ValDataset(Dataset):
    def __init__(self, label_list, transform=None, target_transform=None, loader=default_loader):
        imgs = []
        for index, row in label_list.iterrows():
            imgs.append((row['img_path'], row['labels']))
        self.imgs = imgs
        self.transform = transform
        self.target_transform = target_transform
        self.loader = loader

    def __getitem__(self, index):
        filename, label = self.imgs[index]
        img = self.loader(filename)
        if self.transform is not None:
            img = self.transform(img)
        return img, label

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

# 测试集图片读取
class TestDataset(Dataset):
    def __init__(self, label_list, transform=None, target_transform=None, loader=default_loader):
        imgs = []
        for index, row in label_list.iterrows():
            imgs.append((row['img_path'], row['labels']))
        self.imgs = imgs
        self.transform = transform
        self.target_transform = target_transform
        self.loader = loader

    def __getitem__(self, index):
        filename, label = self.imgs[index]
        img = self.loader(filename)
        if self.transform is not None:
            img = self.transform(img)
        return img, label

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

# 数据增强：在给定角度中随机进行旋转
class FixedRotation(object):
    def __init__(self, angles):
        self.angles = angles

    def __call__(self, img):
        return fixed_rotate(img, self.angles)

def fixed_rotate(img, angles):
    angles = list(angles)
    angles_num = len(angles)
    index = random.randint(0, angles_num - 1)
    return img.rotate(angles[index])

# 训练函数
def train(train_loader, model, criterion, optimizer, epoch):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    acc = AverageMeter()

    # switch to train mode
    model.train()

    end = time.time()
    for i, (images, target) in enumerate(train_loader):
        # measure data loading
        data_time.update(time.time() - end)

        result = []
        for temp in target:
            temp1 = temp.replace("'", "")
            temp2 = temp1.replace(",", "")
            temp3 = temp2.replace(" ", "")
            temp3 = [int(m) for m in temp3]
            for j in range(len(temp3)):
                if temp3[j] == 1:
                    l = j
            result.append(l)
        result1 = np.array(result)
        result2 = torch.from_numpy(result1)
        # print(result2)
        image_var = torch.tensor(images).cuda(async=True)
        label = torch.tensor(result2).cuda(async=True)

        # compute y_pred
        y_pred = model(image_var)
        loss = criterion(y_pred, label)

        # measure accuracy and record loss
        prec, PRED_COUNT = accuracy(y_pred.data, result2, topk=(1, 1))
        losses.update(loss.item(), images.size(0))
        acc.update(prec, PRED_COUNT)

        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % print_freq == 0:
            print('Epoch: [{0}][{1}/{2}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Accuray {acc.val:.3f} ({acc.avg:.3f})'.format(
                epoch, i, len(train_loader), batch_time=batch_time, data_time=data_time, loss=losses, acc=acc))

# 验证函数
def validate(val_loader, model, criterion):
    batch_time = AverageMeter()
    losses = AverageMeter()
    acc = AverageMeter()

    # switch to evaluate mode
    model.eval()

    end = time.time()
    for i, (images, target) in enumerate(val_loader):
        result = []
        for temp in target:
            temp1 = temp.replace("'", "")
            temp2 = temp1.replace(",", "")
            temp3 = temp2.replace(" ", "")
            temp3 = [int(m) for m in temp3]
            for j in range(len(temp3)):
                if temp3[j] == 1:
                    l = j
            result.append(l)
        result1 = np.array(result)
        result2 = torch.from_numpy(result1)
        image_var = torch.tensor(images).cuda(async=True)
        target = torch.tensor(result2).cuda(async=True)

        # compute y_pred
        with torch.no_grad():
            y_pred = model(image_var)
            loss = criterion(y_pred, target)

        # measure accuracy and record loss
        prec, PRED_COUNT = accuracy(y_pred.data, result2, topk=(1, 1))
        losses.update(loss.item(), images.size(0))
        acc.update(prec, PRED_COUNT)

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % print_freq == 0:
            print('Val: [{0}/{1}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Accuray {acc.val:.3f} ({acc.avg:.3f})'.format(
                i, len(val_loader), batch_time=batch_time, loss=losses, acc=acc))

    print(' * Accuray {acc.avg:.3f}'.format(acc=acc), '(Previous Best Acc: %.3f)' % best_precision,
          ' * Loss {loss.avg:.3f}'.format(loss=losses), 'Previous Lowest Loss: %.3f)' % lowest_loss)
    # with open('./result/%s.txt' % file_name, 'a') as file:
    #     file.write('VAIL:Precision: %.8f, Loss: %.8f\n' % (acc.avg, losses.avg))
    return acc.avg, losses.avg

# 测试函数
def test(test_loader, model, criterion):
    batch_time = AverageMeter()
    losses = AverageMeter()
    acc = AverageMeter()
    end = time.time()
    for i, (images, target) in enumerate(test_loader):
        result = []
        for temp in target:
            temp1 = temp.replace("'", "")
            temp2 = temp1.replace(",", "")
            temp3 = temp2.replace(" ", "")
            temp3 = [int(m) for m in temp3]
            for j in range(len(temp3)):
                if temp3[j] == 1:
                    l = j
            result.append(l)
        result1 = np.array(result)
        result2 = torch.from_numpy(result1)
        image_var = torch.tensor(images).cuda(async=True)
        target = torch.tensor(result2).cuda(async=True)

        # compute y_pred
        with torch.no_grad():
            y_pred = model(image_var)
            loss = criterion(y_pred, target)

        # measure accuracy and record loss
        prec, PRED_COUNT = accuracy(y_pred.data, result2, topk=(1, 1))

        losses.update(loss.item(), images.size(0))
        acc.update(prec, PRED_COUNT)
        batch_time.update(time.time() - end)
        end = time.time()
        if i % print_freq == 0:
            print('Test: [{0}/{1}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Accuray {acc.val:.3f} ({acc.avg:.3f})'.format(
                i, len(test_loader), batch_time=batch_time, loss=losses, acc=acc))
        loss_avg = losses.avg
        pre_avg = acc.avg
    print(' * Accuray {acc.avg:.3f}'.format(acc=acc),
          ' * Loss {loss.avg:.3f}'.format(loss=losses))
    with open('./result/%s.txt' % file_name, 'a') as file:
        file.write('TEST:Precision: %.8f, Loss: %.8f\n' % (pre_avg, loss_avg))

# 保存最新模型以及最优模型
def save_checkpoint(state, is_best, is_lowest_loss, filename='./model/%s/checkpoint.pth.tar' % file_name):
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, './model/%s/model_best.pth.tar' % file_name)
    if is_lowest_loss:
        shutil.copyfile(filename, './model/%s/lowest_loss.pth.tar' % file_name)

# 用于计算精度和时间的变化
class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

# 学习率衰减：lr = lr / lr_decay
def adjust_learning_rate():

    nonlocal lr
    lr = lr / lr_decay
    return optim.Adam(model.parameters(), lr, weight_decay=weight_decay, amsgrad=True)

# 计算top K准确率
def accuracy(y_pred, y_actual, topk=(1,)):

    """Computes the precision@k for the specified values of k"""
    final_acc = 0
    maxk = max(topk)
    # for prob_threshold in np.arange(0, 1, 0.01):
    PRED_COUNT = y_actual.size(0)
    PRED_CORRECT_COUNT = 0
    prob, pred = y_pred.topk(maxk, 1, True, True)
    # prob = np.where(prob > prob_threshold, prob, 0)
    for j in range(pred.size(0)):
        if int(y_actual[j]) == int(pred[j]):
            PRED_CORRECT_COUNT += 1
    if PRED_COUNT == 0:
        final_acc = 0
    else:
        final_acc = PRED_CORRECT_COUNT / PRED_COUNT
    return final_acc * 100, PRED_COUNT

# 程序主体

# 设定GPU ID
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
# 小数据集上，batch size不易过大，如出现out of memory，应再调小batch size
batch_size = 64
# 进程数量，最好不要超过电脑最大进程数，尽量能被batch size整除，windows下报错可以改为workers=0
workers = 12

# epoch数量，分stage进行，跑完一个stage后降低学习率进入下一个stage
stage_epochs = [20, 10, 10]
# 初始学习率
lr = 1e-4
# 学习率衰减系数 (new_lr = lr / lr_decay)
lr_decay = 5
# 正则化系数
weight_decay = 1e-4

# 参数初始化
stage = 0
start_epoch = 0
total_epochs = sum(stage_epochs)
best_precision = 0
lowest_loss = 100

# 设定打印频率，即多少step打印一次，用于观察loss和acc的实时变化
# 打印结果中，括号前面为实时loss和acc，括号内部为epoch内平均loss和acc
print_freq = 1
# 验证集比例
val_ratio = 0.12
# 是否只验证，不训练
evaluate = False
# 是否从断点继续跑
resume = False
# 创建inception_v4模型
# model = DInet.net(num_classes=23)
# model = inception_v3.net(num_classes=16)
model = models.inception_v3(pretrained=True)
model.aux_logit = False

in_features = model.fc.in_features

model.fc = nn.Linear(in_features, 16)

# print(model)

# torch.save(model, './model.pkl')
# import pickle

# f = open('dict_word.pkl', 'rb')
# for line in f:
#     print(line)
# word = pickle.load(open("./model.pkl", 'rb'), encoding='utf-8')
#    train = pickle.load(open("./model.pkl", 'rb'),encoding='iso-8859-1')
#    print(train)
#    file=open("./model.txt",'r')
#    for temp in train:
#        file.write(temp)

# model.save("./model.txt")

# model = torch.nn.DataParallel(model).cuda()
model.cuda()
# optionally resume from a checkpoint
if resume:
    checkpoint_path = './model/%s/checkpoint.pth.tar' % file_name
    if os.path.isfile(checkpoint_path):
        print("=> loading checkpoint '{}'".format(checkpoint_path))
        checkpoint = torch.load(checkpoint_path)
        start_epoch = checkpoint['epoch'] + 1
        best_precision = checkpoint['best_precision']
        lowest_loss = checkpoint['lowest_loss']
        stage = checkpoint['stage']
        lr = checkpoint['lr']
        model.load_state_dict(checkpoint['state_dict'])
        # 如果中断点恰好为转换stage的点，需要特殊处理
        if start_epoch in np.cumsum(stage_epochs)[:-1]:
            stage += 1

            optimizer = adjust_learning_rate()
            model.load_state_dict(torch.load('./model/%s/model_best.pth.tar' % file_name)['state_dict'])
        print("=> loaded checkpoint (epoch {})".format(checkpoint['epoch']))
    else:
        print("=> no checkpoint found at '{}'".format(resume))

# 读取训练图片列表
all_data = pd.read_csv('/home/lab326/songpeng/train_val/train.csv')
# 分离训练集和测试集，stratify参数用于分层抽样
train_data_list, val_data_list = train_test_split(all_data, test_size=0.1, stratify=all_data['labels'])
test_data_list = pd.read_csv('/home/lab326/songpeng/train_val/val.csv')
# 读取测试图片列表
# test_data_list = pd.read_csv('data/test.csv')

# 图片归一化，由于采用ImageNet预训练网络，因此这里直接采用ImageNet网络的参数
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

# 训练集图片变换，输入网络的尺寸为384*384
train_data = TrainDataset(train_data_list,
                          transform=transforms.Compose([
                              transforms.Resize((299, 299)),
                              # transforms.ColorJitter(0.15, 0.15, 0.15, 0.075),
                              # transforms.RandomHorizontalFlip(),
                              # transforms.RandomGrayscale(),
                              # # transforms.RandomRotation(20),
                              # FixedRotation([0, 90, 180, 270]),
                              # transforms.RandomCrop(224),
                              transforms.ToTensor(),
                              normalize,
                          ]))

# 验证集图片变换
val_data = ValDataset(val_data_list,
                      transform=transforms.Compose([
                          transforms.Resize((299, 299)),
                          # transforms.RandomHorizontalFlip(),
                          # transforms.RandomRotation(20),
                          # transforms.CenterCrop(224),
                          transforms.ToTensor(),
                          normalize,
                      ]))

# 测试集图片变换
test_data = TestDataset(test_data_list,
                        transform=transforms.Compose([
                            transforms.Resize((299, 299)),
                            # transforms.RandomHorizontalFlip(),
                            # transforms.RandomRotation(20),
                            # transforms.CenterCrop(224),
                            transforms.ToTensor(),
                            normalize,
                        ]))

# 生成图片迭代器
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=workers)
val_loader = DataLoader(val_data, batch_size=batch_size * 2, shuffle=False, pin_memory=False, num_workers=workers)
test_loader = DataLoader(test_data, batch_size=batch_size * 2, shuffle=False, pin_memory=False, num_workers=workers)

# 使用交叉熵损失函数
criterion = nn.CrossEntropyLoss().cuda()

# 优化器，使用带amsgrad的Adam
optimizer = optim.Adam(model.parameters(), lr, weight_decay=weight_decay, amsgrad=True)
Loss_list = []
Accuracy_list = []
if evaluate:
    validate(val_loader, model, criterion)
else:
    # 开始训练
    for epoch in range(start_epoch, total_epochs):
        # train for one epoch
        train(train_loader, model, criterion, optimizer, epoch)
        # evaluate on validation set
        precision, avg_loss = validate(val_loader, model, criterion)
        # precision, avg_loss = validate(val_loader, model, criterion)
        Loss_list.append(avg_loss)
        precision = precision / 100
        Accuracy_list.append(precision)
        # 在日志文件中记录每个epoch的精度和loss
        with open('./result/%s.txt' % file_name, 'a') as acc_file:
            acc_file.write('Epoch: %2d, Precision: %.8f, Loss: %.8f\n' % (epoch, precision, avg_loss))

        # 记录最高精度与最低loss，保存最新模型与最佳模型
        is_best = precision > best_precision
        is_lowest_loss = avg_loss < lowest_loss
        best_precision = max(precision, best_precision)
        lowest_loss = min(avg_loss, lowest_loss)
        state = {
            'epoch': epoch,
            'state_dict': model.state_dict(),
            'best_precision': best_precision,
            'lowest_loss': lowest_loss,
            'stage': stage,
            'lr': lr,
        }
        save_checkpoint(state, is_best, is_lowest_loss)

        # 判断是否进行下一个stage
        if (epoch + 1) in np.cumsum(stage_epochs)[:-1]:
            stage += 1
            lr = lr / lr_decay
            # optim.Adam(model.parameters(), lr, weight_decay=weight_decay, amsgrad=True)

            optimizer = adjust_learning_rate()
            model.load_state_dict(torch.load('./model/%s/model_best.pth.tar' % file_name)['state_dict'])
            print('Step into next stage')
            with open('./result/%s.txt' % file_name, 'a') as acc_file:
                acc_file.write('---------------Step into next stage----------------\n')

# 记录线下最佳分数
with open('./result/%s.txt' % file_name, 'a') as acc_file:
    acc_file.write('* best acc: %.8f  %s\n' % (best_precision, os.path.basename(__file__)))
with open('./result/best_acc.txt', 'a') as acc_file:
    acc_file.write('%s  * best acc: %.8f  %s\n' % (
        time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time())), best_precision,
        os.path.basename(__file__)))

# 读取最佳模型，预测测试集
best_model = torch.load('./model/%s/lowest_loss.pth.tar' % file_name)
model.load_state_dict(best_model['state_dict'])
test(test_loader, model, criterion)

# 释放GPU缓存
torch.cuda.empty_cache()
return Loss_list, Accuracy_list