train

import os
import argparse

import numpy as np
from tqdm import tqdm
import pandas as pd
import joblib
from sklearn.model_selection import StratifiedKFold

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import torch.backends.cudnn as cudnn
from torchvision import transforms

from lib.dataset import Dataset
from lib.utils.utils import AverageMeter, str2bool, RandomErase
from lib.metrics import compute_accuracy
from lib.losses import FocalLoss, LabelSmoothingLoss
from lib.optimizers import RAdam
from lib.models.RA import RA
from lib.models.model_factory import get_model
from lib.models.gcn import SoftLabelGCN
# 新增

import torch
import cv2
from PIL import Image
from sklearn import metrics
from torch.optim.optimizer import Optimizer
import matplotlib.pyplot as plt
import copy

"""The num of channels the encoder will output 
this will be based on the encoder model 
we are using. For VGG16 it's 512"""
encode_out_num_channels = 2048


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument('--name',
                        default=None,
                        help='model name: (default: arch+timestamp)')
    parser.add_argument('--arch',
                        '-a',
                        metavar='ARCH',
                        default='vgg16',
                        help='model architecture: ' + ' (default: resnet34)')
    parser.add_argument('--freeze_bn', default=True, type=str2bool)
    parser.add_argument('--dropout_p', default=0, type=float)
    parser.add_argument('--loss',
                        default='CrossEntropyLoss',
                        choices=[
                            'CrossEntropyLoss', 'FocalLoss', 'MSELoss',
                            'LabelSmoothingLoss'
                        ])
    parser.add_argument('--reg_coef', default=1.0, type=float)
    parser.add_argument('--cls_coef', default=0.1, type=float)
    parser.add_argument('--epochs',
                        default=10,
                        type=int,
                        metavar='N',
                        help='number of total epochs to run')
    parser.add_argument('-b',
                        '--batch_size',
                        default=32,
                        type=int,
                        metavar='N',
                        help='mini-batch size (default: 32)')
    parser.add_argument('--img_size',
                        default=256,
                        type=int,
                        help='input image size (default: 256)')
    parser.add_argument('--input_size',
                        default=224,
                        type=int,
                        help='input image size (default: 224)')
    parser.add_argument('--optimizer', default='Adam')
    parser.add_argument('--pred_type',
                        default='classification',
                        choices=['classification', 'regression'])
    parser.add_argument('--scheduler',
                        default='CosineAnnealingLR',
                        choices=['CosineAnnealingLR', 'ReduceLROnPlateau'])
    parser.add_argument('--lr',
                        '--learning_rate',
                        default=0.0001,
                        type=float,
                        metavar='LR',
                        help='initial learning rate')
    parser.add_argument('--min_lr',
                        default=1e-5,
                        type=float,
                        help='minimum learning rate')
    parser.add_argument('--factor', default=0.5, type=float)
    parser.add_argument('--patience', default=5, type=int)
    parser.add_argument('--momentum', default=0.9, type=float, help='momentum')
    parser.add_argument('--weight_decay',
                        default=1e-4,
                        type=float,
                        help='weight decay')
    parser.add_argument('--nesterov',
                        default=False,
                        type=str2bool,
                        help='nesterov')
    parser.add_argument('--gpus', default='0', type=str)
    parser.add_argument('--mode',
                        default='arcnet',
                        choices=['baseline', 'arcnet', 'gcn'])
    parser.add_argument('--mod',
                        default='aid20',
                        choices=['aid50', 'optail80', 'gcn'])
    # lstm
    parser.add_argument('--lstm_layers', default=3, type=int)
    parser.add_argument('--lstm_hidden', default=256, type=int)
    parser.add_argument('--lstm_recurrence', default=10, type=int)

    # preprocessing
    parser.add_argument('--scale_radius', default=True, type=str2bool)
    parser.add_argument('--normalize', default=False, type=str2bool)
    parser.add_argument('--padding', default=False, type=str2bool)

    # data augmentation
    parser.add_argument('--rotate', default=True, type=str2bool)
    parser.add_argument('--rotate_min', default=-180, type=int)
    parser.add_argument('--rotate_max', default=180, type=int)
    parser.add_argument('--rescale', default=True, type=str2bool)
    parser.add_argument('--rescale_min', default=0.8889, type=float)
    parser.add_argument('--rescale_max', default=1.0, type=float)
    parser.add_argument('--shear', default=True, type=str2bool)
    parser.add_argument('--shear_min', default=-36, type=int)
    parser.add_argument('--shear_max', default=36, type=int)
    parser.add_argument('--translate', default=False, type=str2bool)
    parser.add_argument('--translate_min', default=0, type=float)
    parser.add_argument('--translate_max', default=0, type=float)
    parser.add_argument('--flip', default=True, type=str2bool)
    parser.add_argument('--contrast', default=True, type=str2bool)
    parser.add_argument('--contrast_min', default=0.9, type=float)
    parser.add_argument('--contrast_max', default=1.1, type=float)
    parser.add_argument('--random_erase', default=True, type=str2bool)
    parser.add_argument('--random_erase_prob', default=0.5, type=float)
    parser.add_argument('--random_erase_sl', default=0.02, type=float)
    parser.add_argument('--random_erase_sh', default=0.4, type=float)
    parser.add_argument('--random_erase_r', default=0.3, type=float)

    # dataset
    parser.add_argument('--train_dataset',
                        default='AID',
                        choices=['Images', 'OPTIMAL-31', 'whurs', 'nwpu', 'AID'])
    parser.add_argument('--cv', default=True, type=str2bool)
    parser.add_argument('--n_splits', default=5, type=int)  # 5

    parser.add_argument('--pretrained_model')

    args = parser.parse_args(args=[])

    return args


# senet模块
class SELayer(nn.Module):
    def __init__(self, channel, reduction=16):
        super(SELayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        print(y.shape)
        y = self.fc(y).view(b, c, 1, 1)
        print(y.shape)
        return x * y.expand_as(x)


# 模型
from torch.autograd import Variable
import torch.nn.functional as F
from torchvision import models
import pretrainedmodels
from efficientnet_pytorch import EfficientNet
from lib.models.MobileNetV2 import mobilenet_v2

# 建立一个全局变量的字典，将特征图放在其中
feature_map = {}


# 构建hook函数
def forward_hook1(module, inp, outp):
    feature_map['features3'] = outp


def forward_hook2(module, inp, outp):
    feature_map['features15'] = outp


def forward_hook3(module, inp, outp):
    feature_map['features25'] = outp


def forward_hook4(module, inp, outp):
    feature_map['features29'] = outp


class ChannelAttention(nn.Module):  # 通道注意力
    def __init__(self, in_planes, ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.fc1 = nn.Conv2d(in_planes, in_planes // 16, 1, bias=False)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Conv2d(in_planes // 16, in_planes, 1, bias=False)

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
        out = avg_out + max_out
        return self.sigmoid(out)


class SpatialAttention(nn.Module):  # 空间注意力
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()

        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()
        # self.register_buffer()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)

        max_out, _ = torch.max(x, dim=1, keepdim=True)

        x = torch.cat([avg_out, max_out], dim=1)

        x = self.conv1(x)

        return self.sigmoid(x)


# 训练模型
class RA(nn.Module):
    def __init__(self, cnn_model_name='resnet18', input_size=224, hidden_size=256, layer_num=3, recurrent_num=5,
                 class_num=5, dropout_p=0.5, pretrain=False):
        super(RA, self).__init__()

        self.cnn_model_name = cnn_model_name
        self.cnn = get_pretrained_cnn(
            cnn_model_name, class_num, pretrain)

        self.hidden_size = hidden_size
        self.layer_num = layer_num
        self.recurrent_num = recurrent_num

        self.class_num = class_num
        self.se = SELayer(512, 3)

        self.lstm1 = nn.LSTM(64, self.hidden_size,
                             self.layer_num, batch_first=True)
        self.lstm2 = nn.LSTM(128, self.hidden_size,
                             self.layer_num, batch_first=True)
        self.lstm3 = nn.LSTM(256, self.hidden_size,
                             self.layer_num, batch_first=True)
        self.lstm4 = nn.LSTM(512, self.hidden_size,
                             self.layer_num, batch_first=True)

        self.att1 = nn.Sequential(
            nn.Dropout(p=dropout_p),
            nn.Linear(self.hidden_size, 50176, bias=False),  # 64
        )
        self.att2 = nn.Sequential(
            nn.Dropout(p=dropout_p),
            nn.Linear(self.hidden_size, 12544, bias=False),  # 128
        )
        self.att3 = nn.Sequential(
            nn.Dropout(p=dropout_p),
            nn.Linear(self.hidden_size, 3136, bias=False),  # 256
        )
        self.att4 = nn.Sequential(
            nn.Dropout(p=dropout_p),
            nn.Linear(self.hidden_size, 784, bias=False),  # 512
        )

        self.fc = nn.Sequential(
            nn.Dropout(p=dropout_p),
            nn.Linear(self.hidden_size, self.class_num, bias=False),
        )

        self.tanh = nn.Tanh()
        self.sigmoid = nn.Sigmoid()
        # 新增
        self.c_attention = ChannelAttention(in_planes=512)
        self.s_attention = SpatialAttention()
        self.fd = nn.Sequential(
            nn.Dropout(p=dropout_p),
            nn.Flatten(),
            nn.Linear(512 * 7 * 7, self.class_num, bias=False),
        )
        self.f_last = nn.Sequential(
            nn.Dropout(p=dropout_p),
        )
        self.avgpool = nn.AdaptiveAvgPool2d(1)

    def single_lstm(self, input, h, c, m, flag):
        # pdb.set_trace()

        input = (m.unsqueeze(1) * input).sum(2)
        input = input.unsqueeze(1)  # 加一维

        if flag == 1:
            a, (h, c) = self.lstm1(input, (h, c))  # 返回隐藏状态输出及ct
            # a, (h, c) = self.lstm3(a, (h, c))
            a, (h, c) = self.lstm3(a, (h, c))
            _, (h, c) = self.lstm3(a, (h, c))
        elif flag == 2:
            a, (h, c) = self.lstm3(input, (h, c))
            # a, (h, c) = self.lstm3(a, (h, c))
            # a, (h, c) = self.lstm3(a, (h, c))
            a, (h, c) = self.lstm3(a, (h, c))
            _, (h, c) = self.lstm3(a, (h, c))
        elif flag == 3:
            a, (h, c) = self.lstm4(input, (h, c))  # 输出都是256
            # a, (h, c) = self.lstm3(a, (h, c))
            # a, (h, c) = self.lstm3(a, (h, c))
            a, (h, c) = self.lstm3(a, (h, c))
            _, (h, c) = self.lstm3(a, (h, c))
        elif flag == 4:
            a, (h, c) = self.lstm4(input, (h, c))
            # a, (h, c) = self.lstm3(a, (h, c))
            # a, (h, c) = self.lstm3(a, (h, c))
            a, (h, c) = self.lstm3(a, (h, c))
            _, (h, c) = self.lstm3(a, (h, c))

        output = self.tanh(h[-1])

        if flag == 1:
            m = self.att1(h[-1])
        elif flag == 2:
            m = self.att3(h[-1])
        elif flag == 3:
            m = self.att4(h[-1])
        elif flag == 4:
            m = self.att4(h[-1])

        # self.m = F.softmax(self.m, dim=1)
        m = self.sigmoid(m)

        output = self.fc(output)
        # output = F.softmax(output, dim=1)

        return output

    def forward(self, input):
        if 'efficient' in self.cnn_model_name:
            cnn_x = self.cnn.extract_features(input)  # 1,c,7,7
        elif 'vgg' in self.cnn_model_name:
            cnn_x = self.cnn.features(input)  # 32,512,7,7
        else:
            cnn_x = self.cnn.features(input)  # 1,c,7,7
        # 通道注意力和空间注意力
        # print(cnn_x)
        a = self.c_attention(cnn_x) #[32, 512, 1, 1]
        b = self.s_attention(a)
        c = cnn_x * b

        su = self.fd(c)
        # VGG16的不同层
        features = list(self.cnn.children())[0]

        #         hook_layer1 = features[3]
        #         hook_layer1.register_forward_hook(forward_hook1)

        hook_layer2 = features[15]  # 3-3
        hook_layer2.register_forward_hook(forward_hook2)

        hook_layer3 = features[25]  # 5-1
        hook_layer3.register_forward_hook(forward_hook3)

        hook_layer4 = features[29]  # 5-3
        hook_layer4.register_forward_hook(forward_hook4)

        with torch.no_grad():
            score = self.cnn.features(input)

        c2 = feature_map['features15']
        c2 = self.avgpool(c2)

        c3 = feature_map['features25']
        c3 = self.avgpool(c3)

        c4 = feature_map['features29']
        c4 = self.avgpool(c4)
       
        #         #第一个LSTM
        self.h2 = Variable(torch.zeros(
            self.layer_num, c2.size(0), self.hidden_size)).cuda()
        self.c2 = Variable(torch.zeros(
            self.layer_num, c2.size(0), self.hidden_size)).cuda()
        self.m2 = Variable(torch.ones(c2.size(0), 12544)).cuda()
        # self.m = F.softmax(self.m, dim=1)
        c2 = c2.view(c2.size(0), c2.size(1), -1)  # 将前面的数据展平
        flag = 2
        k2 = self.single_lstm(c2, self.h2, self.c2, self.m2, flag)

        # 第二个LSTM
        self.h3 = Variable(torch.zeros(
            self.layer_num, c3.size(0), self.hidden_size)).cuda()
        self.c3 = Variable(torch.zeros(
            self.layer_num, c3.size(0), self.hidden_size)).cuda()
        self.m3 = Variable(torch.ones(c3.size(0), 3136)).cuda()
        # self.m = F.softmax(self.m, dim=1)
        c3 = c3.view(c3.size(0), c3.size(1), -1)  # 将前面的数据展平

        flag = 3
        k3 = self.single_lstm(c3, self.h3, self.c3, self.m3, flag)

        # 第三LSTM
        self.h4 = Variable(torch.zeros(
            self.layer_num, c4.size(0), self.hidden_size)).cuda()
        self.c4 = Variable(torch.zeros(
            self.layer_num, c4.size(0), self.hidden_size)).cuda()
        self.m4 = Variable(torch.ones(c4.size(0), 784)).cuda()
        # self.m = F.softmax(self.m, dim=1)
        c4 = c4.view(c4.size(0), c4.size(1), -1)  # 将前面的数据展平
        flag = 4
        k4 = self.single_lstm(c4, self.h4, self.c4, self.m4, flag)

        kz = k2 + k3 + k4  # 四个lstm的和
        tr = torch.cat([su, kz], dim=1)

        return F.softmax(tr, dim=1)

    def get_config_optim(self, lr_cnn, lr_lstm):
        return [{'params': self.cnn.parameters(), 'lr': lr_cnn},
                {'params': self.att.parameters(), 'lr': lr_lstm},
                {'params': self.fc.parameters(), 'lr': lr_lstm},
                {'params': self.lstm.parameters(), 'lr': lr_lstm},
                ]


def get_pretrained_cnn(model_name='resnet18', output_num=None, pretrained=True, freeze_bn=False, dropout_p=0, **kwargs):
    if 'efficientnet' in model_name:
        model = EfficientNet.from_pretrained(
            model_name, num_classes=output_num)
        channel = model._fc.in_features

    elif 'densenet' in model_name:
        model = models.__dict__[model_name](
            num_classes=1000, pretrained=pretrained)
        channel = model.classifier.in_features
        model.classifier = nn.Linear(channel, output_num)

    elif 'mobilenet' in model_name:
        model = mobilenet_v2(pretrained=pretrained)
        channel = model.classifier.in_features
        model.classifier = nn.Linear(channel, output_num)

    elif 'vgg16' in model_name:
        pretrained = 'imagenet' if pretrained else None
        model = models.__dict__[model_name](num_classes=1000,
                                            pretrained=pretrained)
        # model.classifier._modules['6'] = nn.Linear(4096, output_num)
        model.classifier = nn.Sequential( nn.Dropout(p=dropout_p),
                                          nn.Linear(73920, output_num, bias=False))
    else:
        pretrained = 'imagenet' if pretrained else None
        model = pretrainedmodels.__dict__[model_name](
            num_classes=1000, pretrained=pretrained)

        if 'dpn' in model_name:
            channel = model.last_linear.in_channels
            model.last_linear = nn.Conv2d(
                channel, output_num, kernel_size=1, bias=True)
        else:
            if 'resnet' in model_name:
                model.avgpool = nn.AdaptiveAvgPool2d(1)
            else:
                model.avg_pool = nn.AdaptiveAvgPool2d(1)
            channel = model.last_linear.in_features
            if dropout_p == 0:
                model.last_linear = nn.Linear(channel, output_num)
            else:
                model.last_linear = nn.Sequential(
                    nn.Dropout(p=dropout_p),
                    nn.Linear(channel, output_num),
                )
    #     if 'resnet' in model_name or 'se_resnext' in model_name:
    #         reduction = 32
    #     elif 'vgg' in model_name:
    #         reduction = 32
    #         channel = 64
    #     else:
    #         reduction = 32
    return model


# 读取数据集
def img_path_generator(dataset='ucm'):
    #     if not os.path.exists('UCMerced_LandUse'):
    #         os.makedirs('UCMerced_LandUse')
    img_dir = '/media/lab713/软件安装/cjs/CJSNet-mast/datasets/%s/' % dataset
    img_path = []
    img_labels = []
    dicts = os.listdir(img_dir)
    for root, _, files in os.walk(img_dir):
        for name in files:
            img_path.append(os.path.join(root, name))
            label_name = root.split('/')[-1]
            img_labels.append(int(dicts.index(label_name)))
    return np.array(img_path, dtype=object), np.array(img_labels), len(dicts)


def train(args, train_loader, model, criterion, optimizer, epoch):
    losses = AverageMeter()
    ac_scores = AverageMeter()

    model.train()



    for i, (input, target) in tqdm(enumerate(train_loader),
                                   total=len(train_loader)):
        input = input.cuda()
        target = target.cuda()

        output = model(input)
        print(output)
        print(target)
        if args.mode == 'baseline':
            output = output
        elif args.mode == 'gcn':
            output, adj = output
        else:
            output = output  # 默认选这个

        if args.pred_type == 'classification':  # 默认选这个
            print(output.shape)
            print(target.long().shape)
            loss = criterion(output, target.long())
        elif args.pred_type == 'regression':
            loss = criterion(output.view(-1), target.float())
        elif args.pred_type == 'multitask':
            loss = args.reg_coef * criterion['regression'](output[:, 0], target.float()) + \
                   args.cls_coef * \
                   criterion['classification'](output[:, 1:], target)
            output = output[:, 0].unsqueeze(1)

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

        ac_score = compute_accuracy(output, target)
        losses.update(loss.item(), input.size(0))
        ac_scores.update(ac_score, input.size(0))
    if args.mode == 'gcn':
        print(torch.max(adj))


    return losses.avg, ac_scores.avg


def validate(args, val_loader, model, criterion):
    losses = AverageMeter()  # 管理变量的更新
    ac_scores = AverageMeter()

    # switch to evaluate mode
    model.eval()

    with torch.no_grad():
        for i, (input, target) in tqdm(enumerate(val_loader),  # 显示进度条的
                                       total=len(val_loader)):
            input = input.cuda()
            target = target.cuda()

            output = model(input)
            if args.mode == 'baseline':
                output = output
            elif args.mode == 'gcn':
                output, adj = output
            else:
                output = output

            if args.pred_type == 'classification':
                loss = criterion(output, target.long())
            elif args.pred_type == 'regression':
                loss = criterion(output.view(-1), target.float())
            elif args.pred_type == 'multitask':
                loss = args.reg_coef * criterion['regression'](output[:, 0], target.float()) + \
                       args.cls_coef * \
                       criterion['classification'](output[:, 1:], target)
                output = output[:, 0].unsqueeze(1)

            ac_score = compute_accuracy(output, target)

            losses.update(loss.item(), input.size(0))
            ac_scores.update(ac_score, input.size(0))

    return losses.avg, ac_scores.avg


# 绘制混淆矩阵

def plot_confusion_matrix(cm, class_names, normalize=True, show_text=True, show_fpfn=False, num=0):
    '''
    Parameters
    ----------
    cm : a nxn dim numpy array.
    class_names: a list of class names (str type)
    normalize: whether to normalize the values
    show_text: whether to show value in each block of the matrix, If matrix is large like 10x10 or 20x20 it's better to set it to false
               because it'll be difficult to read values but you can see the network behaviour via color map.
    show_fpfn: whether to show false positives on GT axis and false negatives on Pred axis. FN -> not detected & FP -> wrong detections
    Returns
    -------
    fig: a plot of confusion matrix along with colorbar
    '''
    if show_fpfn:
        conf_mat = cm
        x_labels = copy.deepcopy(class_names)
        y_labels = copy.deepcopy(class_names)
        x_labels.append('FN')
        y_labels.append('FP')
    else:
        conf_mat = cm[0:cm.shape[0], 0:cm.shape[0]]
        x_labels = class_names
        y_labels = class_names
    my_cmap = 'Blues'  # viridis, seismic, gray, ocean, CMRmap, RdYlBu, rainbow, jet, Blues, Greens, Purples

    c_m = conf_mat

    if normalize:
        row_sums = c_m.sum(axis=1)
        c_m = c_m / row_sums[:, np.newaxis]
        c_m = np.round(c_m, 2)  # 第二个参数设置小数点位数

    print('*' * 80)
    print('NOTE: In confusion_matrix the last coloumn "FP/FN" shows False Positives in Groundtruths \
          \nand False Negatives in Predictions')
    print('*' * 80)

    fig, ax = plt.subplots(figsize=(12, 10))
    im = ax.imshow(c_m, cmap=my_cmap)

    # We want to show all ticks...
    ax.set_xticks(np.arange(len(y_labels)))
    ax.set_yticks(np.arange(len(x_labels)))
    # ... and label them with the respective list entries
    ax.set_xticklabels(y_labels)
    ax.set_yticklabels(x_labels)

    # Rotate the tick labels and set their alignment.
    plt.setp(ax.get_xticklabels(), rotation=45, ha="right", va="top", rotation_mode="anchor")  # ha=right
    plt.xticks(fontsize=13)  # 调整刻度标签字体大小
    plt.yticks(fontsize=13)
    plt.xlabel('xlabel', fontsize=13)  # 调整标签大小
    plt.ylabel('ylabel', fontsize=13)
    if show_text:
        for i in range(len(x_labels)):
            for j in range(len(y_labels)):
                if c_m[i, j] == 0: continue
                if (i != j):
                    text = ax.text(j, i, c_m[i, j], color="k", ha="center", va="center",size=7)  # color=clr_select(i, j)
                else:
                    text = ax.text(j, i, c_m[i, j], color="w", ha="center", va="center",size=7)  # color=clr_select(i, j)
    #     ax.set_title("Normalized Confusion Matrix with UCM Dataset")
    fig.suptitle('Normalized Confusion Matrix with AID', fontsize=13)  # 设置标题大小
    fig.tight_layout()
    plt.xlabel('Predicted Labels')
    plt.ylabel('True Labels')
    sm = plt.cm.ScalarMappable(cmap=my_cmap, norm=plt.Normalize(vmin=0, vmax=1))
    sm._A = []
    plt.colorbar(sm)
    if (num == 0):
        plt.savefig("./aid201.svg")
    elif (num == 1):
        plt.savefig("./aid202.svg")
    elif (num == 2):
        plt.savefig("./aid203.svg")
    elif (num == 3):
        plt.savefig("./aid204.svg")
    elif (num == 4):
        plt.savefig("./aid205.svg")
    elif (num == 5):
        plt.savefig("./train6.jpg")
    # plt.show()

    return fig


def main():
    args = parse_args()

    if args.name is None:
        args.name = '%s_%s' % (args.mod, args.arch)
    if not os.path.exists('models/%s' % args.name):
        os.makedirs('models/%s' % args.name)

    print('Config -----')
    for arg in vars(args):
        print('- %s: %s' % (arg, getattr(args, arg)))
    print('------------')

    with open('models/%s/args.txt' % args.name, 'w') as f:
        for arg in vars(args):
            print('- %s: %s' % (arg, getattr(args, arg)), file=f)

    joblib.dump(args, 'models/%s/args.pkl' % args.name)

    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus

    # switch to benchmark model, a little forward results fluctuation, a little fast training
    cudnn.benchmark = True
    # switch to deterministic model, more stable
    # cudnn.deterministic = True

    img_path, img_labels, num_outputs = img_path_generator(
        dataset=args.train_dataset)
    if args.pred_type == 'regression':
        num_outputs = 1

    skf = StratifiedKFold(n_splits=args.n_splits, shuffle=True, random_state=0)  # k折交叉验证
    img_paths = []
    labels = []
    for fold, (train_idx,
               val_idx) in enumerate(skf.split(img_path, img_labels)):
        img_paths.append((img_path[train_idx], img_path[val_idx]))
        labels.append((img_labels[train_idx], img_labels[val_idx]))

    train_transform = []
    train_transform = transforms.Compose([
        transforms.Resize((args.img_size, args.img_size)),
        # transforms.RandomAffine(
        #     degrees=(args.rotate_min, args.rotate_max) if args.rotate else 0,
        #     translate=(args.translate_min, args.translate_max) if args.translate else None,
        #     scale=(args.rescale_min, args.rescale_max) if args.rescale else None,
        #     shear=(args.shear_min, args.shear_max) if args.shear else None,
        # ),
        transforms.RandomCrop(args.input_size, padding=4),
        transforms.RandomHorizontalFlip(),
        # transforms.RandomVerticalFlip(),
        # transforms.ColorJitter(
        #     brightness=0,
        #     contrast=args.contrast,
        #     saturation=0,
        #     hue=0),
        RandomErase(prob=args.random_erase_prob if args.random_erase else 0,
                    sl=args.random_erase_sl,
                    sh=args.random_erase_sh,
                    r=args.random_erase_r),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])

    val_transform = transforms.Compose([
        transforms.Resize((args.img_size, args.img_size)),
        transforms.CenterCrop(args.input_size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])

    if args.loss == 'CrossEntropyLoss':  # 默认选这个
        criterion = nn.CrossEntropyLoss().cuda()
    elif args.loss == 'FocalLoss':
        criterion = FocalLoss().cuda()
    elif args.loss == 'MSELoss':
        criterion = nn.MSELoss().cuda()
    elif args.loss == 'LabelSmoothingLoss':
        criterion = LabelSmoothingLoss(classes=num_outputs,
                                       smoothing=0.8).cuda()
    else:
        raise NotImplementedError

    folds = []
    best_losses = []
    best_ac_scores = []
    best_epochs = []

    train_loss_history = []
    train_acc_history = []
    val_loss_history = []
    val_acc_history = []

    ###使训练比拟小于测试比例改这
    for fold, ((val_img_paths, train_img_paths),
               (val_labels, train_labels)) in enumerate(zip(img_paths,
                                                            labels)):
        print('Fold [%d/%d]' % (fold + 1, len(img_paths)))

        # if os.path.exists('models/%s/model_%d.pth' % (args.name, fold+1)):
        #     log = pd.read_csv('models/%s/log_%d.csv' % (args.name, fold+1))
        #     best_loss, best_ac_score = log.loc[log['val_loss'].values.argmin(
        #     ), ['val_loss', 'val_score', 'val_ac_score']].values
        #     folds.append(str(fold + 1))
        #     best_losses.append(best_loss)
        #     best_ac_scores.append(best_ac_score)
        #     continue

        # train
        train_set = Dataset(train_img_paths,
                            train_labels,
                            transform=train_transform)
        # 数据初始化
        train_loader = torch.utils.data.DataLoader(train_set,
                                                   batch_size=args.batch_size,
                                                   shuffle=True,
                                                   num_workers=0,
                                                   sampler=None)

        val_set = Dataset(val_img_paths, val_labels, transform=val_transform)
        val_loader = torch.utils.data.DataLoader(val_set,
                                                 batch_size=args.batch_size,
                                                 shuffle=False,
                                                 num_workers=0)

        # create model
        if args.mode == 'baseline':
            model = get_model(model_name=args.arch,
                              num_outputs=num_outputs,
                              freeze_bn=args.freeze_bn,
                              dropout_p=args.dropout_p)
        elif args.mode == 'gcn':
            model_path = 'models/%s/model_%d.pth' % ('baseline_' + args.arch, fold + 1)
            if not os.path.exists(model_path):
                print('%s is not exists' % model_path)
                continue
            model = SoftLabelGCN(cnn_model_name=args.arch,
                                 cnn_pretrained=False,
                                 num_outputs=num_outputs)
            pretrained_dict = torch.load(model_path)
            model_dict = model.cnn.state_dict()
            pretrained_dict = {
                k: v
                for k, v in pretrained_dict.items() if k in model_dict
            }
            model_dict.update(pretrained_dict)
            model.cnn.load_state_dict(model_dict)
            for p in model.cnn.parameters():
                p.requires_grad = False
        else:  # 默认用这个
            # model = RA(cnn_model_name=args.arch, input_size=args.input_size, hidden_size=args.lstm_hidden,
            #            layer_num=args.lstm_layers, recurrent_num=args.lstm_recurrence, class_num=num_outputs, pretrain=True)
            model_path = 'models/%s/model_%d.pth' % ('aid20base_' + args.arch,
                                                     fold + 1)
            if not os.path.exists(model_path):
                print('%s is not exists' % model_path)
                continue

            model = RA(cnn_model_name=args.arch,
                       input_size=args.input_size,
                       hidden_size=args.lstm_hidden,
                       layer_num=args.lstm_layers,
                       recurrent_num=args.lstm_recurrence,
                       class_num=num_outputs)
            # print(model)
            # 计算参数量
            total_params = sum(p.numel() for p in model.parameters())
            print(f'{total_params:,} total parameters.')

            pretrained_dict = torch.load(model_path)

            model_dict = model.cnn.state_dict()
            pretrained_dict = {
                k: v
                for k, v in pretrained_dict.items() if k in model_dict
            }

            model_dict.update(pretrained_dict)
            model.cnn.load_state_dict(model_dict)
            for p in model.cnn.parameters():
                p.requires_grad = False

        device = torch.device('cuda')
        if torch.cuda.device_count() > 1:
            model = nn.DataParallel(model)
        model.to(device)
        # 切换到cpu
        #         device = torch.device("cpu")

        #         model.to(device)

        if args.pretrained_model is not None:
            model.load_state_dict(
                torch.load('models/%s/model_%d.pth' %
                           (args.pretrained_model, fold + 1)))

        # print(model)

        if args.optimizer == 'Adam':
            optimizer = optim.Adam(filter(lambda p: p.requires_grad,
                                          model.parameters()),
                                   lr=args.lr)
        elif args.optimizer == 'AdamW':
            optimizer = optim.AdamW(filter(lambda p: p.requires_grad,
                                           model.parameters()),
                                    lr=args.lr)
        elif args.optimizer == 'RAdam':
            optimizer = RAdam(filter(lambda p: p.requires_grad,
                                     model.parameters()),
                              lr=args.lr)
        elif args.optimizer == 'SGD':
            optimizer = optim.SGD(filter(lambda p: p.requires_grad,
                                         model.parameters()),
                                  lr=args.lr,
                                  momentum=args.momentum,
                                  weight_decay=args.weight_decay,
                                  nesterov=args.nesterov)
            # optimizer = optim.SGD(model.get_config_optim(args.lr, args.lr * 10, args.lr * 10), lr=args.lr,
            #                       momentum=args.momentum, weight_decay=args.weight_decay, nesterov=args.nesterov)

        if args.scheduler == 'CosineAnnealingLR':
            scheduler = lr_scheduler.CosineAnnealingLR(optimizer,
                                                       T_max=args.epochs,
                                                       eta_min=args.min_lr)
        elif args.scheduler == 'ReduceLROnPlateau':
            scheduler = lr_scheduler.ReduceLROnPlateau(optimizer,
                                                       factor=args.factor,
                                                       patience=args.patience,
                                                       verbose=1,
                                                       min_lr=args.min_lr)

        log = pd.DataFrame(index=[],
                           columns=[
                               'epoch',
                               'loss',
                               'ac_score',
                               'val_loss',
                               'val_ac_score',
                           ])
        log = {
            'epoch': [],
            'loss': [],
            'ac_score': [],
            'val_loss': [],
            'val_ac_score': [],
        }

        best_loss = float('inf')
        best_ac_score = 0
        best_epoch = 0

        for epoch in range(args.epochs):
            print('Epoch [%d/%d]' % (epoch + 1, args.epochs))



            # train for one epoch
            train_loss, train_ac_score = train(args, train_loader, model,
                                               criterion, optimizer, epoch)

            # evaluate on validation set
            val_loss, val_ac_score = validate(args, val_loader, model,
                                              criterion)

            if args.scheduler == 'CosineAnnealingLR':
                scheduler.step()
            elif args.scheduler == 'ReduceLROnPlateau':
                scheduler.step(val_loss)

            print(
                'loss %.4f - ac_score %.4f - val_loss %.4f - val_ac_score %.4f'
                % (train_loss, train_ac_score, val_loss, val_ac_score))

            log['epoch'].append(epoch)
            log['loss'].append(train_loss)
            log['ac_score'].append(train_ac_score)
            log['val_loss'].append(val_loss)
            log['val_ac_score'].append(val_ac_score)
            #画训练损失图
            train_loss_history.append(train_loss)
            val_loss_history.append(val_loss)
            train_acc_history.append(train_ac_score)
            val_acc_history.append(val_ac_score)
            pd.DataFrame(log).to_csv('models/%s/log_%d.csv' %
                                     (args.name, fold + 1),
                                     index=False)

            if val_ac_score > best_ac_score:
                if args.mode == 'baseline':
                    torch.save(
                        model.state_dict(),
                        'models/%s/model_%d.pth' % (args.name, fold + 1))
                best_loss = val_loss
                best_ac_score = val_ac_score
                best_epoch = epoch
                print("=> saved best model")


        print('val_loss:  %f' % best_loss)
        print('val_ac_score: %f' % best_ac_score)

        folds.append(str(fold + 1))
        best_losses.append(best_loss)
        best_ac_scores.append(best_ac_score)
        best_epochs.append(best_epoch)

        results = pd.DataFrame({
            'fold':
                folds + ['mean'],
            'best_loss':
                best_losses + [np.mean(best_losses)],
            'best_ac_score':
                best_ac_scores + [np.mean(best_ac_scores)],
            'best_epoch':
                best_epochs + [''],
        })

        print(results)
        results.to_csv('models/%s/results.csv' % args.name, index=False)
        torch.cuda.empty_cache()

        if not args.cv:
            break

        plt.plot(train_acc_history,label = "train_loss")
        plt.plot(val_acc_history, label="val_loss")
        plt.legend()
        plt.xlabel("epochs")
        plt.ylabel("acc")
        plt.show()
        # 混淆矩阵
        def confusion_matrix(preds, labels, conf_matrix):
            preds = torch.argmax(preds, 1)
            for p, t in zip(preds, labels):
                conf_matrix[p, t] += 1
            return conf_matrix

        conf_matrix = torch.zeros(30, 30)  # 与实际类数一样
        for i, (data, target) in tqdm(enumerate(val_loader), total=len(val_loader)):
            output = model(data.to(device))
            conf_matrix = confusion_matrix(output, target, conf_matrix)

        cnf_matrix = conf_matrix.numpy()
        # # ucm 21
        # attack_types = ['agricultural', 'airplane', 'baseballdiamond', 'beach', 'buildings', 'chaparral',
        #                 'denseresidential', 'forest', 'freeway', 'golfcourse', 'harbor', 'intersection', 'mediumresidential',
        #                 'mobilehomepark', 'overpass', 'parkinglot', 'river', 'runway', 'sparseresidential', 'storagetanks',
        #                 'tenniscourt']
        # # optimal 31
        # #         attack_types = ['airplane', 'airport', 'baseball_diamond','basketball_court', 'beach','bridge',
        # #                         'chaparral', 'church', 'circular_farmland', 'commercial_area', 'dense_residential', 'desert', 'forest',
        # #                         'freeway', 'golf_course', 'ground_track_field', 'harbor', 'industrial_area', 'intersection', 'island',
        # #                        'lake','meadow','medium_residential','mobile_home_park','mountain','overpass','parking_lot','railway',
        # #                        'rectangular_farmland','roundabout','runway']
        # # NWPU 45
        # #         attack_types = ['airplane', 'airport', 'baseball_diamond','basketball_court', 'beach','bridge',
        # #                     'chaparral', 'church', 'circular_farmland','cloud', 'commercial_area', 'dense_residential', 'desert', 'forest',
        # #                     'freeway', 'golf_course', 'ground_track_field', 'harbor', 'industrial_area', 'intersection', 'island',
        # #                    'lake','meadow','medium_residential','mobile_home_park','mountain','overpass','palace','parking_lot','railway','railway_station',
        # #                      'rectangular_farmland','river','roundabout','runway','sea_ice','ship','snowberg','sparse_residential','stadium','storage_tank',
        # #                         'tennis_court','terrace','thermal_power_station','wetland'
        # #                     ]
        # # AID 30
        attack_types = ['Airport', 'BareLand', 'BaseballField', 'beach', 'Bridge', 'Center', 'Church', 'Commercial',
                        'DenseResidential',
                        'Desert', 'Farmland', 'Forest', 'Industrial', 'Meadow', 'MediumResidential', 'Mountain', 'Park',
                        'Parking',
                        'Playground', 'Pond', 'Port', 'RailwayStation', 'Resort', 'River', 'School',
                        'SparseResidential', 'Square',
                        'Stadium', 'StorageTanks', 'Viaduct']
        cnf_matrix = plot_confusion_matrix(cnf_matrix, attack_types, normalize=True, show_text=True,
                                           show_fpfn=False,
                                           num=fold)


if __name__ == '__main__':
    main()