from models import ResNet, ConvNet import torch.nn as nn import argparse from utils import UcrDataset, UCR_dataloader, AdvDataset import torch.optim as optim import torch.utils.data import os import random import numpy as np import matplotlib.pyplot as plt parser = argparse.ArgumentParser() parser.add_argument('--test', action='store_true', help='') parser.add_argument('--query_one', action='store_true', help='query the probability of target idx sample') parser.add_argument('--idx', type=int, help='the index of test sample ') parser.add_argument('--gpu', type=str, default='0', help='the index of test sample ') parser.add_argument('--channel_last', type=bool, default=True, help='the channel of data is last or not') parser.add_argument('--n_class', type=int, default=3, help='the class number of dataset') parser.add_argument('--epochs', type=int, default=1500, help='number of epochs to train for') parser.add_argument('--e', default=1499, help='number of epochs to train for') parser.add_argument('--lr', type=float, default=0.0002, help='learning rate, default=0.0002') parser.add_argument('--cuda', action='store_false', help='enables cuda') parser.add_argument('--checkpoints_folder', default='model_checkpoints', help='folder to save checkpoints') parser.add_argument('--manualSeed', type=int, help='manual seed') parser.add_argument('--run_tag', default='StarLightCurves', help='tags for the current run') parser.add_argument('--model', default='f', help='the model type(ResNet,FCN)') parser.add_argument('--normalize', action='store_true', help='') parser.add_argument('--checkpoint_every', default=5, help='number of epochs after which saving checkpoints') opt = parser.parse_args() print(opt) # configure cuda if torch.cuda.is_available() and not opt.cuda: os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpu print("You have a cuda device, so you might want to run with --cuda as option") device = torch.device("cuda:0" if opt.cuda else "cpu") if opt.manualSeed is None: opt.manualSeed = random.randint(1, 10000) print("Random Seed: ", opt.manualSeed) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) def train(l,e): os.makedirs(opt.checkpoints_folder, exist_ok=True) os.makedirs('%s/%s' % (opt.checkpoints_folder, opt.run_tag), exist_ok=True) #这里是获得的train数据 dataset_path = 'data/' + opt.run_tag + '/' + opt.run_tag + '_TRAIN.txt' dataset = UcrDataset(dataset_path, channel_last=opt.channel_last, normalize=True) #这里是result里面的攻击时间序列数据，对抗训练 # attacked_data_path = 'final_result/' + opt.run_tag + '/' + 'attack_time_series.txt' #攻击的时间序列 # attacked_dataset = AdvDataset(txt_file=attacked_data_path) #batch数值小于等于16（这里是10） batch_size = int(min(len(dataset) / 10, 16)) #batch_size = 64 print('dataset length: ', len(dataset)) # print('number of adv examples：', len(attacked_dataset)) print('batch size：', batch_size) #这里batch数值是10 dataloader = UCR_dataloader(dataset, batch_size) # adv_dataloader = UCR_dataloader(attacked_dataset, batch_size) seq_len = dataset.get_seq_len() n_class = opt.n_class print('sequence len:', seq_len) if opt.model == 'r': net = ResNet(n_in=seq_len, n_classes=n_class).to(device) if opt.model == 'f': net = ConvNet(n_in=seq_len, n_classes=n_class).to(device) net.train() criterion = nn.CrossEntropyLoss().to(device) optimizer = optim.Adam(net.parameters(), lr=opt.lr) print('#############！！ Start to Train！！ ###############') #opt.epochs=100(原来等于1499，别忘了改回去） for epoch in range(opt.epochs): for i, (data, label) in enumerate(dataloader): if data.size(0) != batch_size: break data = data.float() data = data.to(device) label = label.long() label = label.to(device) optimizer.zero_grad() output = net(data) loss = criterion(output, label.view(label.size(0))) loss.backward() optimizer.step() print('#######！！！Train！！！########') print('[%d/%d][%d/%d] Loss: %.4f ' % (epoch, opt.epochs, i + 1, len(dataloader), loss.item())) if (epoch % (opt.checkpoint_every * 10) == 0) or (epoch == (opt.epochs - 1)): print('Saving the %dth epoch model.....' % epoch) #训练完保存模型 torch.save(net, '%s/%s/pre_%s%depoch.pth' % (opt.checkpoints_folder, opt.run_tag, opt.model, epoch)) if (epoch % 10 == 0): l.append(loss.item()) e.append(epoch) def te(): data_path = 'data/' + opt.run_tag + '/' + opt.run_tag + '_TEST.txt' dataset = UcrDataset(data_path, channel_last=opt.channel_last, normalize=opt.normalize) batch_size = int(min(len(dataset) / 10, 16)) print('dataset length: ', len(dataset)) print('batch_size:', batch_size) dataloader = UCR_dataloader(dataset, batch_size) type = opt.model model_path = 'model_checkpoints/' + opt.run_tag + '/pre_' + type + str(opt.e) + 'epoch.pth' #model_path ='model_checkpoints/' + opt.run_tag + '/pre_fTrained.pth' model = torch.load(model_path, map_location='cuda:0') with torch.no_grad(): model.eval() total = 0 correct = 0 for i, (data, label) in enumerate(dataloader): data = data.float() data = data.to(device) label = label.long() label = label.to(device) label = label.view(label.size(0)) total += label.size(0) out = model(data).cuda() softmax = nn.Softmax(dim=-1) prob = softmax(out) pred_label = torch.argmax(prob, dim=1) correct += (pred_label == label).sum().item() print('The TEST Accuracy of %s is : %.2f %%' % (data_path, correct / total * 100)) def query_one(idx): data_path = 'data/' + opt.run_tag + '/' + opt.run_tag + '_TEST.txt' test_data = np.loadtxt(data_path) test_data = torch.from_numpy(test_data) test_one = test_data[idx] X = test_one[1:].float() X = X.to(device) y = test_one[0].long() - 1 y = y.to(device) if y < 0: y = opt.n_class - 1 print('ground truth：', y) type = opt.model model_path = 'model_checkpoints/' + opt.run_tag + '/pre_' + type + str(opt.e) + 'epoch.pth' #model_path = 'model_checkpoints/' + opt.run_tag + '/pre_fTrained.pth' model = torch.load(model_path, map_location='cpu') model.eval() out = model(X) softmax = nn.Softmax(dim=-1) prob_vector = softmax(out) print('prob vector：', prob_vector) prob = prob_vector.view(opt.n_class)[y].item() print('Confidence in true class of the %d sample is %.4f ' % (idx, prob)) def plot1(model,loss,epoch): plt.title('The loss of traing '+opt.run_tag+'model', fontstyle='italic') plt.figure(figsize=(6, 4)) plt.plot(epoch,loss, color='b', label='loss') plt.xlabel('epoch', fontsize=12) plt.legend(loc='upper right', fontsize=8) #plt.savefig('loss' + model+'.png') plt.show() if __name__ == '__main__': # if opt.test: # ce() # elif opt.query_one: # query_one(opt.idx) # else: # train() l = [] e = [] train(l,e) # plot1(opt.model,l,e) te()