preload_data.py

from re import sub
import torch,math
import torch.nn as nn
import torchvision.models as models
import sys
sys.path.append('../josh_484u/hand/')
from handDataset import *
from torch.utils.data import Dataset,SequentialSampler,DataLoader,SubsetRandomSampler
import sklearn.metrics as sm
from sklearn import preprocessing
from sklearn.decomposition import PCA
import numpy as np
import torch.nn.functional as F
import torch.optim as optim
from tqdm import tqdm
import argparse
import matplotlib.pyplot as plt
from torch.utils.tensorboard import SummaryWriter
from sklearn.metrics import confusion_matrix
import itertools
from itertools import chain
import pickle
import math
import random as rnd
import statistics
import json
from collections import Counter
import ssl
ssl._create_default_https_context = ssl._create_unverified_context

from handDataset import *
from models import *
from datasets import onehandDataset, allhandDataset, addaug, normalize
from utils import get_random_seed, plot_confusion_matrix, reportOnlineAccuracy, testslice, testonline, LabelSmoothingLoss

writer = SummaryWriter(log_dir='hand/logger')


def dataset_generation(args,trainset,testset):
    if trainset==None:
        test_sampler = SubsetRandomSampler(list(range(testset.__len__())))
        test_loader= torch.utils.data.DataLoader(testset, batch_size=args.batch_size,
                                              shuffle=False, num_workers=2,drop_last=False,sampler=test_sampler)
        return test_loader
    train_sampler = SubsetRandomSampler(list(range(trainset.__len__())))
    test_sampler = SubsetRandomSampler(list(range(testset.__len__())))
    train_loader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size,
                                              shuffle=False, num_workers=2,drop_last=True,sampler=train_sampler)
    test_loader= torch.utils.data.DataLoader(testset, batch_size=args.batch_size,
                                              shuffle=False, num_workers=2,drop_last=False,sampler=test_sampler)
    return train_loader,test_loader


def train_teacher(model, train_loader, optimizer, criterion, device,best_acc,regular):
    model.train()
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=10, verbose=True)
    train_loss = 0
    total_samples=0
    running_corrects=0
    for i, (x,y) in enumerate(train_loader):
        optimizer.zero_grad()
        x,y =x.to(device), y.to(device)
        try:
            output,_=model(x)
        except:
            output = model(x)
        loss = criterion(output, y.type(torch.int64))
        # 计算训练集上的准确率
        prediction = torch.argsort(output, dim=-1, descending=True)
        top1_acc = torch.sum((prediction[:, 0:1] == y.unsqueeze(dim=-1)).any(dim=-1).float()).item()
        running_corrects += top1_acc
        total_samples += x.size(0)

        # 计算训练集损失
        if regular:
            # regularization
            lambda_reg=0.00001
            l2_regularization = torch.tensor(0.).to(device)
            for param in model.parameters():
                l2_regularization += torch.norm(param, p=2)
            # 展示一下两个loss的大致范围
            if i==0:print(f'realloss:{loss},regular loss:{l2_regularization*lambda_reg}')
            loss=loss+l2_regularization*lambda_reg
        loss.backward()
        optimizer.step()
        train_loss += loss.detach().item()
    print(f'[ST_model]acc on traindataset is:{round(running_corrects/total_samples,2)}')
    avg_train_loss = train_loss / len(train_loader)

    scheduler.step(avg_train_loss)  # Update the learning rate based on
    return avg_train_loss

def test(model, test_loader, device):
    model.eval()
    running_corrects, running_corrects_5,total_samples=0, 0, 0
    all_preds,all_targets = [], []
    with torch.no_grad():
        for i, batch in enumerate(test_loader):
            x, y = batch
            x = x.to(device)
            y = y.to(device)
            batchsize = np.array(y.shape)[0]
            try:
                pred,_=model(x)
            except:
                pred = model(x)
            prediction = torch.argsort(pred, dim=-1, descending=True)
            top1_acc = torch.sum((prediction[:, 0:1] == y.unsqueeze(dim=-1)).any(dim=-1).float()).item()
            top5_acc = torch.sum((prediction[:, 0:5] == y.unsqueeze(dim=-1)).any(dim=-1).float()).item()
            running_corrects += top1_acc
            running_corrects_5 += top5_acc
            all_preds.extend(np.vstack(prediction[:, 0:1].cpu()))
            all_targets.extend(y.cpu().numpy())
            total_samples += x.size(0)
        epoch_acc=running_corrects/total_samples
    all_preds = np.array(all_preds)
    all_targets = np.array(all_targets)
    cm = confusion_matrix(all_targets, all_preds)
    class_accuracies = cm.diagonal() / cm.sum(axis=1)
    return epoch_acc,cm,class_accuracies

def testPerman(model, testdata, device, datarange,savename):
    model.eval()
    with torch.no_grad():
        x, y = testdata.x_data, testdata.y_data
        x = x.to(device)
        y = y.to(device)
        #预测
        pred= model(x)
        prediction = torch.argsort(pred, dim=-1, descending=True)
    #all_preds为模型的总预测结果
    #all_targets为总的label
    all_preds=np.array((np.vstack(prediction[:, 0:1].cpu())))
    all_preds = all_preds.reshape(all_preds.shape[1],-1)[0]
    all_targets=np.array(y.cpu().numpy())
    #开始划分人
    numofperson=datarange[1]-datarange[0]+1
    personsize = all_preds.shape[0]//numofperson
    perman_acc={}
    for j in range(numofperson):
        perman_acc['id_'+j]=np.mean(all_preds[:,j*personsize:(j+1)*personsize]==all_targets[j*personsize:(j+1)*personsize])
    #保存每个人的acc
    with open('hand/perman/'+savename+'.josn','w+') as f:
        json.dump(perman_acc,f)
    #保存acc的图片
    plt.clf()
    plt.hist(list(perman_acc.values()))
    plt.savefig('hand/perman/'+savename+'_permanacc.png')
    plt.close()
    return list(perman_acc.values())


def train_and_evaluate(teacher_model, student_model,train_loader, test_loader,alltestset, optimizer, criterion, num_epochs, device,savename_th, savename_st,subjectId_test,onlinetype,regular):
    best_acc_th=0
    best_acc_st=0
    learning_st = 3e-4
    temperature = 20.0
    alpha = 0.4
    epoch_acc_lis,train_loss=[],[]
    class_names =  ['Class 0', 'Class 1', 'Class 2', 'Class 3', 'Class 4', 'Class 5','Class 6', 'Class 7', 'Class 8', 'Class 9', 'Class 10']
    for epoch in tqdm(range(num_epochs)):
        # train teachet model
        train_loss = train_teacher(teacher_model, train_loader, optimizer, criterion, device,best_acc_th,regular)
        epoch_acc,cm,class_acc= test(teacher_model, test_loader, device)
        if epoch_acc > best_acc_th:
            best_acc_th = epoch_acc
            torch.save(teacher_model.state_dict(), 'hand/model/'+savename_th+'.pth')
            plot_confusion_matrix(cm, class_names)
            plt.savefig('hand/figure/'+savename_th+'confusion_matrix.png', bbox_inches='tight')
            plt.close()
            # 每当有最好的模型，就print一下
            allacc,accperman=testsliceP(teacher_model,alltestset,device,subjectId_test,probthres=2.2)
            print('###[TH_model]Current allacc_slice is: ',round(allacc,4),'acc perman is: ',accperman)
        writer.add_scalar('Loss/Teacher_train', train_loss, epoch)
        writer.add_scalar('Acc/Teacher_acc', epoch_acc, epoch)
        print(f"[TH_model]Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, epoch_acc is: {epoch_acc:.4f}, best_acc_th is: {best_acc_th:.4f}, class_acc is: {np.round(class_acc,4)}")
    # # 没过10个epoch就跑一下testslice函数
    #     if onlinetype=='slice':
    #         if epoch%10==0:
    #             allacc,accperman=testsliceP(teacher_model,alltestset,device,subjectId_test,probthres=2.2)
    #             print('###[TH_model]Current allacc_slice is: ',round(allacc,2),'acc perman is: ',accperman)

        #train student model

        optimizer = optim.Adam(chain(student_model.parameters(), teacher_model.parameters()),
                        lr=learning_st, weight_decay=1e-6)
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=10, verbose=True)
        teacher_model.train()
        student_model.train()
        train_loss,total_soft_loss,total_hard_loss = 0,0,0
        for _,(inputs, labels) in enumerate(train_loader):
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            teacher_logits = teacher_model(inputs)
            student_logits = student_model(inputs)
            soft_loss = distillation_loss(student_logits, teacher_logits, temperature)
            hard_loss = nn.CrossEntropyLoss()(student_logits,labels.type(torch.int64))
            loss = (1 - alpha) * soft_loss + alpha * hard_loss
            #print('soft loss',soft_loss,'hard loss',hard_loss)
            if torch.isnan(loss) or torch.isinf(loss):
                print("Loss is NaN or Inf. Skipping this batch.")
                continue
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
            total_soft_loss+=soft_loss.item()
            total_hard_loss+=hard_loss.item()
        avg_train_loss = train_loss / len(train_loader)
        scheduler.step(avg_train_loss)  # Update the learning rate based on
        epoch_acc,cm,class_acc= test(student_model,test_loader,device)
        if epoch_acc > best_acc_st:
            best_acc_st = epoch_acc
            torch.save(student_model.state_dict(), 'hand/model/'+savename_st+'.pth')
            plot_confusion_matrix(cm, class_names)
            allacc,accperman=testsliceP(student_model,alltestset,device,subjectId_test)
            print('###[ST_model]Current allacc_slice is: ',round(allacc,2),'acc perman is: ',accperman)
            plt.savefig('hand/figure/'+savename_st+'confusion_matrix.png', bbox_inches='tight')
        print(f"[ST_model]Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, epoch_acc is: {epoch_acc:.4f}, best_acc_st is: {best_acc_st:.4f}, class_acc is: {class_acc}")
        # if onlinetype=='slice':
        #     if epoch%10==0:
        #         allacc,accperman=testsliceP(student_model,alltestset,device,subjectId_test,probthres=2.2)
        #         print('###[ST_model]Current allacc_slice is: ',round(allacc,2),'acc perman is: ',accperman)
    return best_acc_th,best_acc_st

def train_student(teacher_model,student_model,train_loader,test_loader,finaltestset_st,alltestset,finaltestset,num_epochs,device,savename,subjectId_test,subjectId_finaltest):
    # 定义超参数
    num_epochs=500
    # learning_rate太大的话会优化困难
    learning_rate = 5e-3
    temperature = 20
    # 0.091即1：10
    # alpha=0 即全用hard_loss
    alpha = 0.9
    best_acc=0
    epoch_acc_lis,train_loss=[],[]
    # optimizer = optim.Adam(chain(student_model.parameters(), teacher_model.parameters()),
    #                    lr=learning_rate, weight_decay=1e-6)
    optimizer = optim.Adam(student_model.parameters(),
                       lr=learning_rate,weight_decay=1e-5)
    # scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, patience=10, verbose=True)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer,T_max=15, eta_min=1e-5, verbose=True)
    class_names = ['Class 0', 'Class 1', 'Class 2', 'Class 3', 'Class 4', 'Class 5','Class 6', 'Class 7', 'Class 8', 'Class 9', 'Class 10']
    for epoch in tqdm(range(num_epochs)):
        teacher_model.train()
        student_model.train()
        # if epoch==0:
        #     allteacherlogits=[]
        train_loss,total_soft_loss,total_hard_loss = 0,0,0
        for index,(inputs, labels) in enumerate(train_loader):
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            # if epoch==0:
                # teacher的logits必须detach,不然就会使得梯度graph被重复调用
            teacher_logits = teacher_model(inputs).detach()
                # allteacherlogits.append(teacher_logits.unsqueeze(0))
            # else:
            # teacher_logits=allteacherlogits[index,:,:]
            student_logits = student_model(inputs)
            loss,soft_loss,hard_loss = distillation_loss(student_logits, teacher_logits, labels.type(torch.int64),temperature,alpha)
            # print('soft loss',soft_loss,'hard loss',hard_loss)
            if torch.isnan(loss) or torch.isinf(loss):
                print("Loss is NaN or Inf. Skipping this batch.")
                continue
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
            total_soft_loss+=soft_loss.item()
            total_hard_loss+=hard_loss.item()
        # 储存teacher模型的logits
        # if epoch==0:
        #     allteacherlogits=torch.cat(allteacherlogits,dim=0)
        print(f'average soft_loss is {total_soft_loss / len(train_loader)}, while average hard_loss is {total_hard_loss / len(train_loader)}')
        avg_train_loss = train_loss / len(train_loader)
        # scheduler.step(avg_train_loss)  # Update the learning rate based on
        scheduler.step()
        epoch_acc,cm,class_acc= test(student_model,test_loader,device)
        if epoch_acc > best_acc:
            best_acc = epoch_acc
            torch.save(student_model.state_dict(), 'hand/model/'+savename+'.pth')
            plot_confusion_matrix(cm, class_names)
            allacc,accperman=testsliceP(student_model,alltestset,device,subjectId_test)
            print('###Current allacc_slice is: ',round(allacc,4),'acc perman is: ',accperman)
            plt.savefig('hand/figure/'+savename+'confusion_matrix.png', bbox_inches='tight')
            plt.close()
            # 对于validation set上面最优的模型，下面在测试集上试一下
            testepoch_acc,testcm,testclass_acc = test(student_model, finaltestset_st, device)
            print(f'on the test dataset, acc is {testepoch_acc:.4f}, class_acc is {testclass_acc}')
            # 测试一下集成后acc
            allacc,accperman=testsliceP(student_model,finaltestset,device,subjectId_finaltest)
            print('on the test dataset, Current allacc_slice is: ',round(allacc,4),'acc perman is: ',accperman)

        print(f"Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, epoch_acc is: {epoch_acc:.4f}, best_acc is: {best_acc:.4f}, class_acc is: {class_acc}")
    return loss, epoch_acc

def testsliceP(model,alltestset,device,subjectId_test,probthres=2):
    model.eval()
    testbatch=50
    # 概率集成判断
    all_preds=[]
    all_targets=[]
    with torch.no_grad():
        for batch in range(int(alltestset.x_data.shape[0]//testbatch+1)):
            x, y = alltestset.x_data[batch*testbatch:min((batch+1)*testbatch,alltestset.x_data.shape[0]),:,:,:], alltestset.y_data[batch*testbatch:min((batch+1)*testbatch,alltestset.x_data.shape[0])]
            x = x.to(device)
            y = y.to(device)
            #预测
            pred= model(x)
            # 将pred通过softmax转化为概率
            # 应用softmax
            # pred_prob = nn.functional.softmax(pred, dim=1).cpu().numpy()
            pred_prob=pred.cpu().numpy()
            #all_preds为模型的总预测结果
            #all_targets为总的label
            targets=np.array(y.cpu().numpy())
            all_targets.append(targets)
            all_preds.append(pred_prob)
    all_preds=np.concatenate(all_preds,axis=0)
    all_targets=np.concatenate(all_targets,axis=0)
    # for pred in range(all_preds.shape[0]):
    #     finderror=np.argmax(all_preds[pred,:])==all_targets[pred]
    #     if not finderror:
    #         print(f'target is{all_targets[pred]},the predict prob is{all_preds[pred,:]}')
    accperman={}
    #每6个target就投票算一个label
    for i in range(int(all_preds.shape[0]/6)):
        allprob=all_preds[i*6:(i+1)*6,:]
        allprob=np.sum(allprob,axis=0)
        max_class = np.argmax(allprob)
        max_prob = allprob[max_class]
        vote=max_class
        # if max_prob<probthres :
        #     # 0表示休息
        #     vote=0
        # else:
        #     vote=max_class
        # 取一个中间值是target
        targetkey=all_targets[i*6+3]
        subject=subjectId_test[i*6+3]
        if subject not in accperman.keys():
            accperman[subject]=[]
        if targetkey==vote:
            accperman[subject].append(1)
        else:
            accperman[subject].append(0)
    for subject,sublist in accperman.items():
        accperman[subject]=np.mean(sublist)
    allacc=np.mean(list(accperman.values()))
    # record acc of every person
    return allacc, accperman


def distillation_loss(student_logits, teacher_logits,labels,temperature,alpha):
    # soft_teacher_logits = torch.softmax(teacher_logits / temperature, dim=-1)
    # soft_student_logits = torch.softmax(student_logits, dim=-1)
    # loss = torch.mean(torch.sum(-soft_teacher_logits * torch.log(soft_student_logits), dim=-1))
    soft_loss=nn.KLDivLoss(reduction='batchmean')(F.log_softmax(student_logits/temperature, dim=1),
                             F.softmax(teacher_logits/temperature, dim=1))
    hard_loss=F.cross_entropy(student_logits, labels)
    KD_loss = soft_loss * (alpha * temperature * temperature) + hard_loss * (1. - alpha)
    return KD_loss, soft_loss, hard_loss

parser = argparse.ArgumentParser()
parser.add_argument('--device', default='0,1,cpu', type=str, help='set device')
parser.add_argument('--batch_size', default=1, type=int,help='batch_size')
parser.add_argument('--train_loader', default='', help='train data loader')
parser.add_argument('--val_loader', default='', help='val data loader')
parser.add_argument('--test_loader', default='', help='test data loader')
parser.add_argument('--input_size', default=40, type=int,help='batch_size')
parser.add_argument('--feature_dim', default=20, type=int,help='batch_size')
parser.add_argument('--out_features', default=100, type=int,help='batch_size')
parser.add_argument('--num_class', default=3, type=int,help='batch_size')
parser.add_argument('--lr1', default=1e-4, type=float,help='batch_size')
parser.add_argument('--model', default='', help='model')
parser.add_argument('--num_epochs', default=1e2, type=int,help='iterations')
parser.add_argument('--loss_type', default='', type=str,help='loss_type')
parser.add_argument('--keywords', default='', type=str,help='keywords')
parser.add_argument('--validation_split', default=0.1, type=float,help='augmentation_num')
parser.add_argument('--testdata_split', default=0.1, type=float,help='stage1_len')
parser.add_argument('--ROC_save_path', default='./', help='ROC_save_path')
parser.add_argument('--loss_save_path', default='', type=str,help='loss_save_path')
parser.add_argument('--model_save_path', default='', type=str,help='model_save_path')
args = parser.parse_args()
args.device = "cuda:1" if torch.cuda.is_available() else "cpu"
args.num_epochs=500
args.num_classes=11
args.batch_size=100
random_seed= 42
args.lr1=3e-4


def main(mode):
    get_random_seed(random_seed)

    # onlinetype='slice'
    # if onlinetype=='slice':
    #     # args.batch_size=125
    #     args.batch_size=50
    # ## train teacher network
    # device=torch.device(args.device)
    # num_epochs=args.num_epochs     # Define optimizer and loss function
    # loss_func = nn.CrossEntropyLoss(label_smoothing=0.2)    # 定义教师网络 - ResNet
    # # loss_func = LabelSmoothingLoss(epsilon=0.2, num_classes=args.num_classes)
    # # loss_func = nn.CrossEntropyLoss()
    # #subjectId的信息储存到allhandDataset类里面了
    # # alltrainset,alltestset=allhandDataset(1,[[1,20],[31,40]],'train',args.num_classes),allhandDataset(2,[[1,20],[31,40]],'test',args.num_classes)
    # alltrainset,alltestset=allhandDataset(1,[[1,20]],'train',args.num_classes,online=onlinetype,augtype='regular'),allhandDataset(2,[[1,20]],'test',args.num_classes,online=onlinetype,augtype=None)
    # train_loader_th,test_loader_th=dataset_generation(args,alltrainset,alltestset)
    # teacher_model = CustomDNN_slice_test(args.num_classes,feature_dim=256) # 使用预训练的 ResNet50 模型
    # # teacher_model=VGGbnmodel(args.num_classes,feature_dim=512)
    # # resnet50_teacher_days_10fold2_sub2是一个训练到了58%的模型
    # # #  resnet50_teacher_days_10fold3是一个加了0.1位置编码,并训练到了58.44%的模型,在用sgd训练400个epoch可以到59.3%
    # # teacher_model.load_state_dict(torch.load('hand/model/resnet50_teacher_days_10fold3.pth'))
    # teacher_model.to(device)
    # # optimizer_th = optim.Adam(teacher_model.parameters(), lr=args.lr1,weight_decay=1e-6)
    # optimizer_th = optim.Adam(teacher_model.parameters(), lr=args.lr1)
    # # testmodel_teacher_days_10fold是一个79%的acc的模型
    # savename = 'DNN_teacher_days_test14'
    # best_acc, teacher_model = train_and_evaluate(teacher_model, train_loader_th, test_loader_th,alltestset,optimizer_th, loss_func, num_epochs, device, savename,subjectId_test=alltestset.subjectId,onlinetype=onlinetype,regular=True) # 定义学生网络
    # print("best_acc is: ",best_acc)




    # # only for test acc based on existing model
    # ## train teacher network
    # onlinetype='slice'
    # device=torch.device(args.device)
    # # loss_func = nn.CrossEntropyLoss(label_smoothing=0.05)    # 定义教师网络 - ResNet
    # loss_func = nn.CrossEntropyLoss()
    # #subjectId的信息储存到allhandDataset类里面了
    # # alltrainset,alltestset=allhandDataset(1,[[1,20],[31,40]],'train',args.num_classes),allhandDataset(2,[[1,20],[31,40]],'test',args.num_classes)
    # alltestset=allhandDataset(2,[[1,20]],'test',args.num_classes,online=onlinetype)
    # test_loader_th=dataset_generation(args,trainset=None,testset=alltestset)
    # teacher_model = CustomDNN_slice_test(args.num_classes,feature_dim=256) # 使用预训练的 ResNet50 模型
    # # teacher_model=VGGbnmodel(args.num_classes,feature_dim=512)
    # # resnet50_teacher_days_10fold2_sub2是一个训练到了58%的模型
    # # #  resnet50_teacher_days_10fold3是一个加了0.1位置编码,并训练到了58.44%的模型,在用sgd训练400个epoch可以到59.3%
    # teacher_model.load_state_dict(torch.load('hand/model/DNN_teacher_days_test2.pth'))
    # teacher_model.to(device)
    # # optimizer_th = optim.Adam(teacher_model.parameters(), lr=args.lr1,weight_decay=1e-6)
    # # optimizer_th = optim.Adam(teacher_model.parameters(), lr=args.lr1)
    # # testmodel_teacher_days_10fold是一个79%的acc的模型
    # epoch_acc,cm,class_acc= test(teacher_model, test_loader_th, device)
    # allacc,accperman=testsliceP(teacher_model,alltestset,device,alltestset.subjectId,probthres=2.5)
    # print(f'acc is:{epoch_acc},slice acc is:{allacc}')
    # # best_acc, teacher_model = train_and_evaluate(teacher_model, train_loader_th, test_loader_th,alltestset,optimizer_th, loss_func, num_epochs, device, savename,subjectId_test=alltestset.subjectId,onlinetype=onlinetype) # 定义学生网络
    # # print("best_acc is: ",best_acc)



    # onlinetype='original'
    # if onlinetype=='slice':
    #     args.batch_size=125
    # ## train teacher network
    # device=torch.device(args.device)
    # num_epochs=args.num_epochs     # Define optimizer and loss function
    # # loss_func = nn.CrossEntropyLoss(label_smoothing=0.05)    # 定义教师网络 - ResNet
    # loss_func = nn.CrossEntropyLoss()
    # #subjectId的信息储存到allhandDataset类里面了
    # # alltrainset,alltestset=allhandDataset(1,[[1,20],[31,40]],'train',args.num_classes),allhandDataset(2,[[1,20],[31,40]],'test',args.num_classes)
    # alltrainset,alltestset=allhandDataset(1,[[1,20]],'train',args.num_classes,online=onlinetype,aug=True),allhandDataset(2,[[1,20]],'test',args.num_classes,online=onlinetype)
    # train_loader_th,test_loader_th=dataset_generation(args,alltrainset,alltestset)
    # teacher_model = CustomDNN_test5(args.num_classes,feature_dim=256) # 使用预训练的 ResNet50 模型
    # # teacher_model=VGGbnmodel(args.num_classes,feature_dim=512)
    # # resnet50_teacher_days_10fold2_sub2是一个训练到了58%的模型
    # # #  resnet50_teacher_days_10fold3是一个加了0.1位置编码,并训练到了58.44%的模型,在用sgd训练400个epoch可以到59.3%
    # # teacher_model.load_state_dict(torch.load('hand/model/resnet50_teacher_days_10fold3.pth'))
    # teacher_model.to(device)
    # # optimizer_th = optim.Adam(teacher_model.parameters(), lr=args.lr1,weight_decay=1e-6)
    # optimizer_th = optim.Adam(teacher_model.parameters(), lr=args.lr1)
    # # testmodel_teacher_days_10fold是一个79%的acc的模型
    # savename = 'DNN_teacher_days_testAug'
    # best_acc, teacher_model = train_and_evaluate(teacher_model, train_loader_th, test_loader_th,alltestset,optimizer_th, loss_func, num_epochs, device,
    #                                              savename,subjectId_test=alltestset.subjectId,onlinetype=onlinetype,regular=False) # 定义学生网络
    # print("best_acc is: ",best_acc)


    ## train teacher network online
    # args.batch_size=150
    # device=torch.device(args.device)
    # num_epochs=args.num_epochs     # Define optimizer and loss function
    # loss_func = nn.CrossEntropyLoss()    # 定义教师网络 - ResNet50
    # alltrainset,alltestset=allhandDataset(1,[1,85],'train',online=True),allhandDataset(2,[1,85],'test',online= True)
    # train_loader_th,test_loader_th=dataset_generation(args,alltrainset,alltestset)
    # teacher_model = CustomResNet50(args.num_classes) # 使用预训练的 ResNet50 模型
    # # teacher_model.load_state_dict(torch.load('hand/model/resnet50_teacher_days_10fold1.pth'))
    # teacher_model.to(device)
    # optimizer_th = optim.Adam(teacher_model.parameters(), lr=args.lr1,weight_decay=1e-6)
    # savename = 'resnet50_teacher_days_10fold1'
    # best_acc, teacher_model = train_and_evaluate(teacher_model, train_loader_th, test_loader_th, optimizer_th, loss_func, num_epochs, device, savename) # 定义学生网络
    # print("best_acc is: ",best_acc)

    # train student network， not simultaneusly
    if mode=='stu':
        onlinetype='slice'
        if onlinetype=='slice':
            # args.batch_size=125
            args.batch_size=250
        device=torch.device(args.device)
        num_epochs=args.num_epochs     # Define optimizer and loss function
        loss_func = nn.CrossEntropyLoss()    # 定义教师网络 - ResNet50
        # alltrainset,alltestset=allhandDataset(1,[2,94],'train'),allhandDataset(2,[2,94],'test')
        # train_loader_th,test_loader_th=dataset_generation(args,alltrainset,alltestset)
        student_model = CustomLSTM(args.num_classes)
        teacher_model = CustomDNN_slice_test(args.num_classes,feature_dim=256) # 使用预训练的 vgc16 模型
        # stat(teacher_model,(32,32,32))
        teacher_model.load_state_dict(torch.load('hand/model/DNN_teacher_days_2fold1_testalldata.pth'))
        teacher_model.to(device)
        # optimizer_th = optim.Adam(teacher_model.parameters(), lr=args.lr1,weight_decay=1e-6)
        # best_acc, teacher_model = train_and_evaluate(teacher_model, train_loader_th, test_loader_th, optimizer_th, loss_func, num_epochs, device) # 定义学生网络
        # print("best_acc is: ",best_acc)
        # onetrainset,onetestset=allhandDataset(1,[[37,41]],'train',args.num_classes,online=onlinetype,augtype='space'),allhandDataset(2,[[1,4],[6,20],[31,36]],'test',args.num_classes,online=onlinetype,augtype=None)
        onetrainset=torch.load('hand/data_preload/stu_fold1_train.pt')
        onetestset=torch.load('hand/data_preload/th_fold1_test.pt')
        # finaltestset=allhandDataset(2,[[37,41]],'test',args.num_classes,online=onlinetype,augtype=None)
        finaltestset=torch.load('hand/data_preload/stu_fold1_test.pt')
        train_loader_st,test_loader_st=dataset_generation(args,onetrainset,onetestset)
        finaltestset_st=dataset_generation(args,None,finaltestset)
        # student_model = CustomCNN_sml(args.num_classes) # 不使用预训练的 ResNet18 模型
        student_model = CustomLSTM(args.num_classes)
        # 这里可以比较一下两个模型
        # stat(student_model,(32,32,32))
        # student_model = CustomDNN_slice_test(args.num_classes,256)
        # resnet18
        # student_model = SmallResNet(BasicBlock, [2,2,2,2])
        student_model.to(device)    # Train and evaluate model
        savename = 'DNN_student_days_CNN2_2fold1'
        loss, epoch_acc_st = train_student(teacher_model,student_model,train_loader_st,test_loader_st,finaltestset_st,onetestset,finaltestset,num_epochs,device, savename,onetestset.subjectId,finaltestset.subjectId)

    if mode=='splitD':
        onlinetype='slice'
        th_fold1=[[1,24]]
        th_fold2=[[1,18],[25,30]]
        th_fold3=[[1,12],[19,30]]
        th_fold4=[[1,6],[13,30]]
        th_fold5=[[7,30]]

        stu=[[1,30]]
        stu_fold1=[[25,30]]
        stu_fold2=[[19,24]]
        stu_fold3=[[13,18]]
        stu_fold4=[[7,12]]
        stu_fold5=[[1,6]]
        train_namelist=['th_fold1','th_fold2','th_fold3','th_fold4','th_fold5','stu','stu_fold1','stu_fold2','stu_fold3','stu_fold4','stu_fold5']
        trainlist=[th_fold1,th_fold2,th_fold3,th_fold4,th_fold5,stu,stu_fold1,stu_fold2,stu_fold3,stu_fold4,stu_fold5]

        test_namelist=['th_fold1','th_fold2','th_fold3','th_fold4','th_fold5','stu','stu_fold1','stu_fold2','stu_fold3','stu_fold4','stu_fold5']
        testlist=[th_fold1,th_fold2,th_fold3,th_fold4,th_fold5,stu,stu_fold1,stu_fold2,stu_fold3,stu_fold4,stu_fold5]


        for index,sublist in enumerate(trainlist):
            allfiles=os.listdir('hand/data_preload/')
            if train_namelist[index]+'_train.pt' in allfiles:continue
            print(f'processing {train_namelist[index]}_train.pt')
            alltrainset=allhandDataset(1,sublist,'train',args.num_classes,online=onlinetype,augtype='space')
            torch.save(alltrainset,'hand/data_preload/'+train_namelist[index]+'_train.pt')

        for index,sublist in enumerate(testlist):
            allfiles=os.listdir('hand/data_preload/')
            if test_namelist[index]+'_test.pt' in allfiles:continue
            print(f'processing {test_namelist[index]}_test.pt')
            alltestset=allhandDataset(2,sublist,'test',args.num_classes,online=onlinetype,augtype=None)
            torch.save(alltestset,'hand/data_preload/'+test_namelist[index]+'_test.pt')

    if mode=='splitD2':
        onlinetype='slice'

        sublist=[[1,30]]
        subid=[]
        for lis in sublist:
            subid=subid + list(range(lis[0],lis[1]+1))


        for id in subid:
            allfiles=os.listdir('/Data2/Data_users/liujionghui/Data_TS_renumbered/PR_preload/')
            if f'{id}_session1.pt' in allfiles:continue
            print(f'processing {id}_session1.pt')
            # 增强的数据当中，rest的数量不足，干脆不增强了
            alltrainset=allhandDataset(1,[id,id],'train',args.num_classes,online=onlinetype,augtype=None)
            torch.save(alltrainset,'/Data2/Data_users/liujionghui/Data_TS_renumbered/PR_preload/'+f'{id}_session1.pt')
            alltrainset=None

        for id in subid:
            allfiles=os.listdir('/Data2/Data_users/liujionghui/Data_TS_renumbered/PR_preload/')
            if f'{id}_session2.pt' in allfiles:continue
            print(f'processing {id}_session2.pt')
            alltestset=allhandDataset(2,[id,id],'test',args.num_classes,online=onlinetype,augtype=None)
            torch.save(alltestset,'/Data2/Data_users/liujionghui/Data_TS_renumbered/PR_preload/'+f'{id}_session2.pt')
            alltestset=None

    #test student perman
    # device=torch.device(args.device)
    # #定义所用数据的范围
    # datarange_test=[1,2]
    # alltestset=allhandDataset(2,datarange_test,'test',online= True)
    # student_model = CustomLSTM(args.num_classes) # 不使用预训练的 ResNet18 模型
    # student_model.load_state_dict(torch.load('hand/model/resnet50_teacher_days_10fold1.pth'))
    # student_model.to(device)    # Train and evaluate model
    # savename = 'onlinetest_resnet50_teacher_days_10fold1'
    # accuracy=testPerman(student_model, alltestset, device, datarange_test,savename)
    # print('accuracy is:',accuracy)

    # # test teacher model online
    # device=torch.device(args.device)
    # num_epochs=args.num_epochs     # Define optimizer and loss function
    # loss_func = nn.CrossEntropyLoss()    # 定义教师网络 - ResNet50
    # #定义所用数据的范围
    # datarange_test=[1,2]
    # alltestset=allhandDataset(2,datarange_test,'test',online= True)
    # teacher_model = OnlineCustomResNet50(args.num_classes) # 使用预训练的 ResNet50 模型
    # teacher_model.load_state_dict(torch.load('hand/model/resnet50_teacher_days_10fold1.pth'))
    # teacher_model.to(device)
    # savename = 'onlinetest_resnet50_teacher_days_10fold1'
    # accuracy, faults=testonline(teacher_model, alltestset, device, datarange_test,savename)
    # print('accuracy is:',accuracy)
    # print('faults ratio is',faults)

# mode='stu'
mode='splitD2'
main(mode)

""" def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

model = CustomCNN(num_classes=10)
total_params = count_parameters(model)
print(f"Total number of parameters in the network: {total_params}") """