def main(): args = parser.parse_args() torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) np.random.seed(args.seed) if args.remark != None: args.remark = args.remark else: args.remark = args.dataset + "-" + args.task + "-" + args.norm if args.dataset == "shapenet": args.num_class = 16 else: args.num_class = 40 def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) experiment_dir = Path('/data-x/g12/zhangjie/3dIP/exp/v2') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath('pruning') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath(args.remark + "_" + timestr) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = experiment_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG_curve''' title = args.dataset + "-" + args.task + "-" + args.norm + "-" + "Pruning" logger_loss = Logger(os.path.join(log_dir, 'log_loss.txt'), title=title) logger_loss.set_names(['Valid Public Loss', 'Valid Private Loss']) logger_acc = Logger(os.path.join(log_dir, 'log_acc.txt'), title=title) logger_acc.set_names(['Valid Public Acc.', 'Valid Private Acc.']) '''LOG''' #创建log文件 logger = logging.getLogger("Model") #log的名字 logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) #log的最低等级 file_handler.setFormatter(formatter) logger.addHandler(file_handler) #log文件名 log_string('PARAMETER ...') log_string(args) '''DATA LOADING''' log_string('Load pruning test dataset ...') if args.dataset == "shapenet": testDataLoader = getData.get_dataLoader(train=False, Shapenet=True, batchsize=args.batch_size) else: testDataLoader = getData.get_dataLoader(train=False, Shapenet=False, batchsize=args.batch_size) log_string('Load finished ...') '''MODEL LOADING''' num_class = args.num_class MODEL = importlib.import_module(args.model) shutil.copy('./models/%s.py' % args.model, str(experiment_dir)) shutil.copy('./models/pointnet_util.py', str(experiment_dir)) shutil.copy('prun0.py', str(experiment_dir)) shutil.copytree('./models/layers', str(experiment_dir) + "/layers") shutil.copytree('./data', str(experiment_dir) + "/data") shutil.copytree('./utils', str(experiment_dir) + "/utils") classifier = MODEL.get_model(num_class, channel=3).cuda() pprint(classifier) pth_dir = '/data-x/g12/zhangjie/3dIP/exp/v2/classification/' + args.dataset + "-" \ + args.task + "-" + args.norm + "/checkpoints/best_model.pth" log_string('pre-trained model chk pth: %s' % pth_dir) checkpoint = torch.load(pth_dir) model_dict = checkpoint['model_state_dict'] print('Total : {}'.format(len(model_dict))) print("best epoch", checkpoint['epoch']) classifier.load_state_dict(model_dict) classifier.cuda() p_num = get_parameter_number(classifier) log_string('Original trainable parameter: %s' % p_num) '''TESTING ORIGINAL''' logger.info('Test original model...') with torch.no_grad(): _, instance_acc, class_acc = test(classifier, testDataLoader, num_class=args.num_class, ind=0) _, instance_acc2, class_acc2 = test(classifier, testDataLoader, num_class=args.num_class, ind=1) log_string( 'Original Instance Public Accuracy: %f, Class Public Accuracy: %f' % (instance_acc, class_acc)) log_string( 'Original Instance Private Accuracy: %f, Class Private Accuracy: %f' % (instance_acc2, class_acc2)) '''PRUNING''' logger.info('Start testing of pruning...') for perc in [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]: time_start = datetime.datetime.now() classifier.load_state_dict(model_dict) p_num = get_parameter_number(classifier) log_string('Original trainable parameter: %s' % p_num) '''Testing pruning model''' logger.info('Testing pruning model--%d%%' % perc) pruning_net(classifier, perc) classifier.cuda() p_num = get_parameter_number(classifier) log_string('Pruning %02d%% -- trainable parameter: %s' % (perc, p_num)) with torch.no_grad(): val_loss1, test_instance_acc1, class_acc1 = test( classifier, testDataLoader, num_class=args.num_class, ind=0) val_loss2, test_instance_acc2, class_acc2 = test( classifier, testDataLoader, num_class=args.num_class, ind=1) log_string( 'Pruning %02d%%-Test Instance Public Accuracy: %f, Class Public Accuracy: %f' % (perc, test_instance_acc1, class_acc1)) log_string( 'Pruning %02d%%-Test Instance Private Accuracy: %f, Class Private Accuracy: %f' % (perc, test_instance_acc2, class_acc2)) # val_loss = (val_loss1 + val_loss2)/2 # test_instance_acc = (test_instance_acc1 + test_instance_acc2)/2 logger_loss.append([val_loss1, val_loss2]) logger_acc.append([test_instance_acc1, test_instance_acc2]) time_end = datetime.datetime.now() time_span_str = str((time_end - time_start).seconds) log_string('Epoch time : %s S' % (time_span_str)) logger_loss.close() logger_loss.plot_prun() savefig(os.path.join(log_dir, 'log_loss.eps')) logger_acc.close() logger_acc.plot_prun() savefig(os.path.join(log_dir, 'log_acc.eps')) logger.info('End of pruning...')
def main(): args = parser.parse_args() torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) np.random.seed(args.seed) if args.remark != None: args.remark = args.remark else: args.remark = args.dataset + "-" + args.task + "-" + args.norm if args.dataset == "shapenet": args.num_class = 16 else: args.num_class = 40 def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) # experiment_dir = Path('./exp/v1/') experiment_dir = Path('/data-x/g12/zhangjie/3dIP/exp/v1') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath('classification') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath(args.remark + "_" + timestr) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = experiment_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG_curve''' title = args.dataset + "-" + args.task + "-" + args.norm logger_loss = Logger(os.path.join(log_dir, 'log_loss_v1.txt'), title=title) logger_loss.set_names( ['Train Loss', 'Valid Clean Loss', 'Valid Trigger Loss']) logger_acc = Logger(os.path.join(log_dir, 'log_acc_v1.txt'), title=title) logger_acc.set_names( ['Train Acc.', 'Valid Clean Acc.', 'Valid Trigger Acc.']) '''LOG''' #创建log文件 logger = logging.getLogger("Model") #log的名字 logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) #log的最低等级 file_handler.setFormatter(formatter) logger.addHandler(file_handler) #log文件名 log_string('PARAMETER ...') log_string(args) '''DATA LOADING''' log_string('Load dataset ...') if args.dataset == "shapenet": trainDataLoader = getData.get_dataLoader(train=True, Shapenet=True, batchsize=args.batch_size) testDataLoader = getData.get_dataLoader(train=False, Shapenet=True, batchsize=args.batch_size) triggerDataLoader = getData2.get_dataLoader(Shapenet=True, batchsize=args.batch_size) else: trainDataLoader = getData.get_dataLoader(train=True, Shapenet=False, batchsize=args.batch_size) testDataLoader = getData.get_dataLoader(train=False, Shapenet=False, batchsize=args.batch_size) triggerDataLoader = getData2.get_dataLoader(Shapenet=False, batchsize=args.batch_size) wminputs, wmtargets = [], [] for wm_idx, (wminput, wmtarget) in enumerate(triggerDataLoader): wminputs.append(wminput) wmtargets.append(wmtarget) '''MODEL LOADING''' num_class = args.num_class MODEL = importlib.import_module(args.model) shutil.copy('./models/%s.py' % args.model, str(experiment_dir)) shutil.copy('./models/pointnet_util.py', str(experiment_dir)) shutil.copy('train_1_cls.py', str(experiment_dir)) shutil.copy('./data/getData.py', str(experiment_dir)) shutil.copy('./data/getData2.py', str(experiment_dir)) shutil.copytree('./models/layers', str(experiment_dir) + "/layers") classifier = MODEL.get_model(num_class, channel=3).cuda() # classifier = MODEL.get_model(num_class,normal_channel=args.normal).cuda() criterion = MODEL.get_loss().cuda() pprint(classifier) try: checkpoint = torch.load( str(experiment_dir) + '/checkpoints/best_model.pth') start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) log_string('Use pretrain model') except: log_string('No existing model, starting training from scratch...') start_epoch = 0 if args.optimizer == 'Adam': optimizer = torch.optim.Adam(classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate) else: optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7) global_epoch = 0 global_step = 0 best_instance_acc = 0.0 best_class_acc = 0.0 mean_correct = [] mean_loss = [] '''TRANING''' logger.info('Start training...') for epoch in range(start_epoch, args.epoch): time_start = datetime.datetime.now() log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) scheduler.step() for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): points, target = data wm_id = np.random.randint(len(wminputs)) points = torch.cat( [points, wminputs[(wm_id + batch_id) % len(wminputs)]], dim=0) #随机选择wininputs和inputscat target = torch.cat( [target, wmtargets[(wm_id + batch_id) % len(wminputs)]], dim=0) points = points.data.numpy() points = provider.random_point_dropout( points) #provider是自己写的一个对点云操作的函数,随机dropout,置为第一个点的值 points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3]) #点的放缩 points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3]) #点的偏移 points = torch.Tensor(points) points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() optimizer.zero_grad() classifier = classifier.train() pred, trans_feat = classifier(points) loss = criterion(pred, target.long(), trans_feat) pred_choice = pred.data.max(1)[1] correct = pred_choice.eq(target.long().data).cpu().sum() mean_correct.append(correct.item() / float(points.size()[0])) loss.backward() optimizer.step() global_step += 1 mean_loss.append(loss.item() / float(points.size()[0])) train_loss = np.mean(mean_loss) train_instance_acc = np.mean(mean_correct) log_string('Train Instance Accuracy: %f' % train_instance_acc) with torch.no_grad(): val_loss, instance_acc, class_acc = test(classifier, testDataLoader, num_class=args.num_class) val_loss2, instance_acc2, class_acc2 = test( classifier, triggerDataLoader, num_class=args.num_class) if (instance_acc >= best_instance_acc): best_instance_acc = instance_acc best_epoch = epoch + 1 if (class_acc >= best_class_acc): best_class_acc = class_acc log_string('Test Clean Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc)) log_string('Best Clean Instance Accuracy: %f, Class Accuracy: %f' % (best_instance_acc, best_class_acc)) log_string( 'Test Trigger Accuracy: %f, Trigger Class Accuracy: %f' % (instance_acc2, class_acc2)) if (instance_acc >= best_instance_acc): logger.info('Save model...') savepath = str(checkpoints_dir) + '/best_model.pth' log_string('Saving at %s' % savepath) log_string('best_epoch %s' % str(best_epoch)) state = { 'epoch': best_epoch, 'clean instance_acc': instance_acc, 'clean class_acc': class_acc, 'trigger instance_acc': instance_acc2, 'trigger class_acc': class_acc2, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) global_epoch += 1 logger_loss.append([train_loss, val_loss, val_loss2]) logger_acc.append([train_instance_acc, instance_acc, instance_acc2]) time_end = datetime.datetime.now() time_span_str = str((time_end - time_start).seconds) log_string('Epoch time : %s S' % (time_span_str)) logger_loss.close() logger_loss.plot() savefig(os.path.join(log_dir, 'log_loss_v3.eps')) logger_acc.close() logger_acc.plot() savefig(os.path.join(log_dir, 'log_acc_v3.eps')) log_string('best_epoch %s' % str(best_epoch)) logger.info('End of training...')
def main(): args = parser.parse_args() # torch.manual_seed(args.seed) # torch.cuda.manual_seed(args.seed) # np.random.seed(args.seed) if args.remark != None: args.remark = args.remark else: args.remark = args.dataset + "-" + args.task + "-" + args.norm if args.dataset == "shapenet": args.num_class = 16 else: args.num_class = 40 if 'fake2-' in args.type: args.flipperc = 0 print('No Flip') elif 'fake3-' in args.type: args.flipperc = 0 print('No Flip') def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' if args.task == 'baseline': experiment_dir_root = Path('/data-x/g12/zhangjie/3dIP/baseline') else: experiment_dir_root = Path('/data-x/g12/zhangjie/3dIP/ours') experiment_dir_root.mkdir(exist_ok=True) experiment_dir = experiment_dir_root.joinpath('ambiguity_attack') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath(args.remark) experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath(args.type) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = experiment_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' logger = logging.getLogger("Model") #log name logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) #log file name log_string('PARAMETER ...') log_string(args) '''DATA LOADING''' log_string('Load dataset ...') if args.dataset == "shapenet": trainDataLoader = getData.get_dataLoader(train=True, Shapenet=True, batchsize=args.batch_size) testDataLoader = getData.get_dataLoader(train=False, Shapenet=True, batchsize=args.batch_size) else: trainDataLoader = getData.get_dataLoader(train=True, Shapenet=False, batchsize=args.batch_size) testDataLoader = getData.get_dataLoader(train=False, Shapenet=False, batchsize=args.batch_size) log_string('Finished ...') log_string('Load model ...') '''MODEL LOADING''' num_class = args.num_class MODEL = importlib.import_module(args.model) # copy model file to exp dir shutil.copy('./models/%s.py' % args.model, str(experiment_dir)) shutil.copy('./models/pointnet_util.py', str(experiment_dir)) shutil.copy('attack2.py', str(experiment_dir)) classifier = MODEL.get_model(num_class, channel=3).cuda() # criterion = MODEL.get_loss().cuda() criterion = nn.NLLLoss().cuda() sd = experiment_dir_root.joinpath('classification') sd.mkdir(exist_ok=True) sd = sd.joinpath(str(args.remark)) sd.mkdir(exist_ok=True) sd = sd.joinpath('checkpoints/best_model.pth') checkpoint = torch.load(sd) classifier.load_state_dict(checkpoint['model_state_dict']) for param in classifier.parameters(): param.requires_grad_(False) origpassport = [] fakepassport = [] for n, m in classifier.named_modules(): if n in ['convp1', 'convp2', 'convp3', 'p1', 'p2', 'fc3']: key, skey = m.__getattr__('key_private').data.clone( ), m.__getattr__('skey_private').data.clone() origpassport.append(key.cuda()) origpassport.append(skey.cuda()) m.__delattr__('key_private') # 删除属性 m.__delattr__('skey_private') # fake like random onise if 'fake2-' in args.type: # fake random m.register_parameter( 'key_private', nn.Parameter(torch.randn(*key.size()) * 0.001, requires_grad=True)) m.register_parameter( 'skey_private', nn.Parameter(torch.randn(*skey.size()) * 0.001, requires_grad=True)) # fake slightly modify ori else: # fake slightly modify ori m.register_parameter( 'key_private', nn.Parameter(key.clone() + torch.randn(*key.size()) * 0.001, requires_grad=True)) m.register_parameter( 'skey_private', nn.Parameter(skey.clone() + torch.randn(*skey.size()) * 0.001, requires_grad=True)) fakepassport.append(m.__getattr__('key_private')) fakepassport.append(m.__getattr__('skey_private')) if args.task == 'ours': if args.type != 'fake2': for layer in m.fc.modules(): if isinstance(layer, nn.Linear): nn.init.xavier_normal_(layer.weight) for i in m.fc.parameters(): i.requires_grad = True if args.flipperc != 0: log_string(f'Reverse {args.flipperc * 100:.2f}% of binary signature') for name, m in classifier.named_modules(): if name in ['convp1', 'convp2', 'convp3', 'p1', 'p2', 'p3']: mflip = args.flipperc oldb = m.sign_loss_private.b newb = oldb.clone() npidx = np.arange(len(oldb)) # bit 长度 randsize = int(oldb.view(-1).size(0) * mflip) randomidx = np.random.choice(npidx, randsize, replace=False) # 随机选择 newb[randomidx] = oldb[randomidx] * -1 # reverse bit 进行翻转 m.sign_loss_private.set_b(newb) classifier.cuda() optimizer = torch.optim.SGD(fakepassport, lr=args.lr, momentum=0.9, weight_decay=0.0005) # scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7) scheduler = None def run_cs(): cs = [] for d1, d2 in zip(origpassport, fakepassport): d1 = d1.view(d1.size(0), -1) d2 = d2.view(d2.size(0), -1) cs.append(F.cosine_similarity(d1, d2).item()) return cs classifier.train() global_epoch = 0 global_step = 0 best_instance_acc = 0.0 best_class_acc = 0.0 mean_correct2 = [] mean_loss2 = [] start_epoch = 0 mse_criterion = nn.MSELoss() cs_criterion = nn.CosineSimilarity() '''TRANING''' logger.info('Start training...') for epoch in range(start_epoch, args.epoch): time_start = datetime.datetime.now() log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) optimizer.zero_grad() signacc_meter = 0 signloss_meter = 0 if scheduler is not None: scheduler.step() for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): points, target = data points = points.data.numpy() points = provider.random_point_dropout(points) points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3]) points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3]) points = torch.Tensor(points) points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() optimizer.zero_grad() classifier = classifier.train() # loss define pred2, _ = classifier(points, ind=1) loss2 = criterion(pred2, target.long()) mean_loss2.append(loss2.item() / float(points.size()[0])) pred_choice2 = pred2.data.max(1)[1] correct2 = pred_choice2.eq(target.long().data).cpu().sum() mean_correct2.append(correct2.item() / float(points.size()[0])) signacc = torch.tensor(0.).cuda() count = 0 for m in classifier.modules(): if isinstance(m, SignLoss): signacc += m.get_acc() count += 1 try: signacc_meter += signacc.item() / count except: pass sign_loss = torch.tensor(0.).cuda() for m in classifier.modules(): if isinstance(m, SignLoss): sign_loss += m.loss signloss_meter += sign_loss loss = loss2 maximizeloss = torch.tensor(0.).cuda() mseloss = torch.tensor(0.).cuda() csloss = torch.tensor(0.).cuda() for l, r in zip(origpassport, fakepassport): mse = mse_criterion(l, r) cs = cs_criterion(l.view(1, -1), r.view(1, -1)).mean() csloss += cs mseloss += mse maximizeloss += 1 / mse if 'fake2-' in args.type: (loss).backward() # only cross-entropy loss backward fake2 elif 'fake3-' in args.type: (loss + maximizeloss).backward() # csloss do not backward kafe3 else: (loss + maximizeloss + 1000 * sign_loss).backward() # csloss backward #fake3_S # (loss + 1000 * sign_loss).backward() # csloss backward #fake3_S torch.nn.utils.clip_grad_norm_(fakepassport, 2) optimizer.step() global_step += 1 signacc = signacc_meter / len(trainDataLoader) log_string('Train Sign Accuracy: %f' % signacc) signloss = signloss_meter / len(trainDataLoader) log_string('Train Sign Loss: %f' % signloss) train_instance_acc2 = np.mean(mean_correct2) log_string('Train Instance Private Accuracy: %f' % train_instance_acc2) with torch.no_grad(): cs = run_cs() log_string( f'Cosine Similarity of Real and Maximize passport: {sum(cs) / len(origpassport):.4f}' ) val_loss2, test_instance_acc2, class_acc2, singloss2, signacc2 = test( classifier, testDataLoader, num_class=args.num_class, ind=1) log_string( 'Test Instance Private Accuracy: %f, Class Private Accuracy: %f' % (test_instance_acc2, class_acc2)) log_string('Test Private Sign Accuracy: %f' % (signacc2)) test_instance_acc = test_instance_acc2 class_acc = class_acc2 if (test_instance_acc >= best_instance_acc): best_instance_acc = test_instance_acc best_epoch = epoch + 1 if (class_acc >= best_class_acc): best_class_acc = class_acc log_string( 'Test Instance Average Accuracy: %f, Class Average Accuracy: %f' % (test_instance_acc, class_acc)) log_string( 'Best Instance Average Accuracy: %f, Class Average Accuracy: %f' % (best_instance_acc, best_class_acc)) if (test_instance_acc >= best_instance_acc): logger.info('Save model...') savepath = str(checkpoints_dir) + '/best_attack_model.pth' log_string('Saving at %s' % savepath) log_string('best_epoch %s' % str(best_epoch)) state = { 'epoch': best_epoch, 'instance_acc': test_instance_acc, 'class_acc': class_acc, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'origpassport': origpassport, 'fakepassport': fakepassport } torch.save(state, savepath) global_epoch += 1 time_end = datetime.datetime.now() time_span_str = str((time_end - time_start).seconds) log_string('Epoch time : %s S' % (time_span_str)) log_string('best_epoch %s' % str(best_epoch)) logger.info('End of training...')
def main(args): def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) experiment_dir = Path('./exp/') # experiment_dir = Path('/data-x/g12/zhangjie/pointnet/exp/') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath('classification') experiment_dir.mkdir(exist_ok=True) if args.log_dir is None: experiment_dir = experiment_dir.joinpath(timestr) else: experiment_dir = experiment_dir.joinpath(args.log_dir) experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath(args.remark + "_" + timestr) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = experiment_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' #创建log文件 args = parse_args() logger = logging.getLogger("Model") #log的名字 # print("FFFFFFFF",logger) #<Logger Model (WARNING)> logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) #log的最低等级 file_handler.setFormatter(formatter) logger.addHandler(file_handler) #log文件名 log_string('PARAMETER ...') # print("FFFFFFF",logger.info) #<bound method Logger.info of <Logger Model (INFO)>> log_string(args) '''DATA LOADING''' log_string('Load dataset ...') # # DATA_PATH = 'data/modelnet40_normal_resampled/' # DATA_PATH = 'data/modelnet40_normal_resampled/' # # TRAIN_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='train', # normal_channel=args.normal) # TEST_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='test', # normal_channel=args.normal) # trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size, shuffle=True, num_workers=4) # testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=4) trainDataLoader = getData.get_dataLoader(train=True) testDataLoader = getData.get_dataLoader(train=False) '''MODEL LOADING''' num_class = args.num_class MODEL = importlib.import_module(args.model) # 当在写代码时,我们希望能够根据传入的选项设置,如args.model来确定要导入使用的是哪个model.py文件,而不是一股脑地导入, 这种时候就需要用上python的动态导入模块 # 复制model文件到exp——dir shutil.copy('./models/%s.py' % args.model, str(experiment_dir)) shutil.copy('./models/pointnet_util.py', str(experiment_dir)) classifier = MODEL.get_model(num_class, normal_channel=args.normal).cuda() criterion = MODEL.get_loss().cuda() pprint(classifier) try: checkpoint = torch.load( str(experiment_dir) + '/checkpoints/best_model.pth') start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) log_string('Use pretrain model') except: log_string('No existing model, starting training from scratch...') start_epoch = 0 if args.optimizer == 'Adam': optimizer = torch.optim.Adam(classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate) else: optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7) global_epoch = 0 global_step = 0 best_instance_acc = 0.0 best_class_acc = 0.0 mean_correct = [] '''TRANING''' logger.info('Start training...') for epoch in range(start_epoch, args.epoch): log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) scheduler.step() for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): points, target = data # print("POINTS",points.shape) [24Batch,1024,6] # print("TARGET",target.shape) [24Batch,1] points = points.data.numpy() points = provider.random_point_dropout( points) #provider是自己写的一个对点云操作的函数,随机dropout,置为第一个点的值 points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3]) #点的放缩 points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3]) #点的偏移 points = torch.Tensor(points) # target = target[:, 0] points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() optimizer.zero_grad() classifier = classifier.train() pred, trans_feat = classifier(points) loss = criterion(pred, target.long(), trans_feat) pred_choice = pred.data.max(1)[1] correct = pred_choice.eq(target.long().data).cpu().sum() mean_correct.append(correct.item() / float(points.size()[0])) loss.backward() optimizer.step() global_step += 1 train_instance_acc = np.mean(mean_correct) log_string('Train Instance Accuracy: %f' % train_instance_acc) with torch.no_grad(): instance_acc, class_acc = test(classifier.eval(), testDataLoader, num_class=args.num_class) if (instance_acc >= best_instance_acc): best_instance_acc = instance_acc best_epoch = epoch + 1 if (class_acc >= best_class_acc): best_class_acc = class_acc log_string('Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc)) log_string('Best Instance Accuracy: %f, Class Accuracy: %f' % (best_instance_acc, best_class_acc)) if (instance_acc >= best_instance_acc): logger.info('Save model...') savepath = str(checkpoints_dir) + '/best_model.pth' log_string('Saving at %s' % savepath) state = { 'epoch': best_epoch, 'instance_acc': instance_acc, 'class_acc': class_acc, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) global_epoch += 1 logger.info('End of training...')
def main(): args = parser.parse_args() torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) np.random.seed(args.seed) if args.remark != None: args.remark = args.remark else: args.remark = args.dataset + "-" + args.task + "-" + args.norm if args.dataset == "shapenet": args.num_class = 16 else: args.num_class = 40 def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) # experiment_dir = Path('./exp/v3/') if args.task == 'baseline': experiment_dir = Path('/data-x/g12/zhangjie/3dIP/baseline') else: experiment_dir = Path('/data-x/g12/zhangjie/3dIP/ours') experiment_dir = experiment_dir.joinpath('classification') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath(args.remark + "_" + timestr) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = experiment_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG_curve''' title = args.dataset + "-" + args.task + "-" + args.norm logger_loss = Logger(os.path.join(log_dir, 'log_loss_v3.txt'), title=title) logger_loss.set_names([ 'Train Pub&Pri Loss', 'Train Public Loss', 'Train Private Loss', 'Valid Pub-Clean loss', 'Valid Pub-Trigger Loss', 'Valid Pri-Clean Loss', 'Valid Pri-Trigger Loss', ]) logger_acc = Logger(os.path.join(log_dir, 'log_acc_v3.txt'), title=title) logger_acc.set_names([ 'Train Pub-Combine Acc.', 'Valid Pub-Clean Acc.', 'Valid Pub-Trigger Acc.', 'Train Pri-Combine Acc.', 'Valid Pri-Clean Acc.', 'Valid Pri-Trigger Acc.' ]) '''LOG''' #创建log文件 logger = logging.getLogger("Model") #log的名字 logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) #log的最低等级 file_handler.setFormatter(formatter) logger.addHandler(file_handler) #log文件名 log_string('PARAMETER ...') log_string(args) '''DATA LOADING''' log_string('Load dataset ...') if args.dataset == "shapenet": trainDataLoader = getData.get_dataLoader(train=True, Shapenet=True, batchsize=args.batch_size) testDataLoader = getData.get_dataLoader(train=False, Shapenet=True, batchsize=args.batch_size) triggerDataLoader = getData2.get_dataLoader(Shapenet=True, T1=args.T1, batchsize=args.batch_size) else: trainDataLoader = getData.get_dataLoader(train=True, Shapenet=False, batchsize=args.batch_size) testDataLoader = getData.get_dataLoader(train=False, Shapenet=False, batchsize=args.batch_size) triggerDataLoader = getData2.get_dataLoader(Shapenet=False, T1=args.T1, batchsize=args.batch_size) wminputs, wmtargets = [], [] for wm_idx, (wminput, wmtarget) in enumerate(triggerDataLoader): wminputs.append(wminput) wmtargets.append(wmtarget) '''MODEL LOADING''' num_class = args.num_class MODEL = importlib.import_module(args.model) shutil.copy('./models/%s.py' % args.model, str(experiment_dir)) shutil.copy('./models/pointnet_util.py', str(experiment_dir)) shutil.copy('train_3_cls.py', str(experiment_dir)) shutil.copy('./data/getData.py', str(experiment_dir)) shutil.copy('./data/getData2.py', str(experiment_dir)) shutil.copytree('./models/layers', str(experiment_dir) + "/layers") classifier = MODEL.get_model(num_class, channel=3).cuda() criterion = MODEL.get_loss().cuda() pprint(classifier) try: checkpoint = torch.load( str(experiment_dir) + '/checkpoints/best_model.pth') start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) log_string('Use pretrain model') except: log_string('No existing model, starting training from scratch...') start_epoch = 0 if args.optimizer == 'Adam': optimizer = torch.optim.Adam(classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate) else: optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7) global_epoch = 0 global_step = 0 best_instance_acc = 0.0 best_class_acc = 0.0 mean_correct = [] mean_correct2 = [] mean_loss = [] mean_loss1 = [] mean_loss2 = [] '''TRANING''' logger.info('Start training...') for epoch in range(start_epoch, args.epoch): time_start = datetime.datetime.now() log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) scheduler.step() wm_id = np.random.randint(len(wminputs)) for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): points, target = data points = torch.cat( [points, wminputs[(wm_id + batch_id) % len(wminputs)]], dim=0) #随机选择wininputs和inputscat target = torch.cat( [target, wmtargets[(wm_id + batch_id) % len(wminputs)]], dim=0) points = points.data.numpy() points = provider.random_point_dropout( points) #provider是自己写的一个对点云操作的函数,随机dropout,置为第一个点的值 points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3]) #点的放缩 points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3]) #点的偏移 points = torch.Tensor(points) points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() optimizer.zero_grad() classifier = classifier.train() for m in classifier.modules(): if isinstance(m, SignLoss): m.reset() loss1 = torch.tensor(0.).cuda() loss2 = torch.tensor(0.).cuda() sign_loss = torch.tensor(0.).cuda() for ind in range(2): if ind == 0: pred, trans_feat = classifier(points, ind=ind) loss1 = criterion(pred, target.long(), trans_feat) mean_loss1.append(loss1.item() / float(points.size()[0])) pred_choice = pred.data.max(1)[1] correct = pred_choice.eq(target.long().data).cpu().sum() mean_correct.append(correct.item() / float(points.size()[0])) else: pred2, trans_feat2 = classifier(points, ind=ind) loss2 = criterion(pred2, target.long(), trans_feat2) mean_loss2.append(loss2.item() / float(points.size()[0])) pred_choice2 = pred2.data.max(1)[1] correct2 = pred_choice2.eq(target.long().data).cpu().sum() mean_correct2.append(correct2.item() / float(points.size()[0])) for m in classifier.modules(): if isinstance(m, SignLoss): sign_loss += m.loss loss = args.beta * loss1 + loss2 + sign_loss mean_loss.append(loss.item() / float(points.size()[0])) loss.backward() optimizer.step() global_step += 1 train_instance_acc = np.mean(mean_correct) train_instance_acc2 = np.mean(mean_correct2) train_loss = np.mean(mean_loss) train_loss1 = np.mean(mean_loss1) train_loss2 = np.mean(mean_loss2) log_string('Train Combine Public Accuracy: %f' % train_instance_acc) log_string('Train Combine Private Accuracy: %f' % train_instance_acc2) sign_acc = torch.tensor(0.).cuda() count = 0 for m in classifier.modules(): if isinstance(m, SignLoss): sign_acc += m.acc count += 1 if count != 0: sign_acc /= count log_string('Sign Accuracy: %f' % sign_acc) res = {} avg_private = 0 count_private = 0 with torch.no_grad(): if args.task == 'ours': for name, m in classifier.named_modules(): if name in [ 'convp1', 'convp2', 'convp3', 'p1', 'p2', 'p3' ]: signbit, _ = m.get_scale(ind=1) signbit = signbit.view(-1).sign() privatebit = m.b detection = ( signbit == privatebit).float().mean().item() res['private_' + name] = detection avg_private += detection count_private += 1 elif args.task == 'baseline': for name, m in classifier.named_modules(): if name in [ 'convp1', 'convp2', 'convp3', 'p1', 'p2', 'p3' ]: signbit = m.get_scale(ind=1).view(-1).sign() privatebit = m.b detection = ( signbit == privatebit).float().mean().item() res['private_' + name] = detection avg_private += detection count_private += 1 log_string('Private Sign Detection Accuracy: %f' % (avg_private / count_private * 100)) for ind in range(2): if ind == 0: val_loss1, test_instance_acc1, class_acc1 = test( classifier, testDataLoader, num_class=args.num_class, ind=0) val_loss_wm1, instance_acc_wm, class_acc_wm = test( classifier, triggerDataLoader, num_class=args.num_class, ind=0) else: val_loss2, test_instance_acc2, class_acc2 = test( classifier, testDataLoader, num_class=args.num_class, ind=1) val_loss_wm2, instance_acc_wm2, class_acc_wm2 = test( classifier, triggerDataLoader, num_class=args.num_class, ind=1) log_string( 'Test Clean Public Accuracy: %f, Class Public Accuracy: %f' % (test_instance_acc1, class_acc1)) log_string( 'Test Clean Private Accuracy: %f, Class Private Accuracy: %f' % (test_instance_acc2, class_acc2)) log_string( 'Test Trigger Public Accuracy: %f, Trigger Class Public Accuracy: %f' % (instance_acc_wm, class_acc_wm)) log_string( 'Test Trigger Private Accuracy: %f, Trigger Class Private Accuracy: %f' % (instance_acc_wm2, class_acc_wm2)) test_instance_acc = (test_instance_acc1 + test_instance_acc2) / 2 class_acc = (class_acc1 + class_acc2) / 2 if (test_instance_acc >= best_instance_acc): best_instance_acc = test_instance_acc best_epoch = epoch + 1 if (class_acc >= best_class_acc): best_class_acc = class_acc log_string( 'Test Combine Average Accuracy: %f, Class Average Accuracy: %f' % (test_instance_acc, class_acc)) log_string( 'Best Combine Average Accuracy: %f, Class Average Accuracy: %f' % (best_instance_acc, best_class_acc)) if (test_instance_acc >= best_instance_acc): logger.info('Save model...') savepath = str(checkpoints_dir) + '/best_model.pth' log_string('Saving at %s' % savepath) log_string('best_epoch %s' % str(best_epoch)) state = { 'epoch': best_epoch, 'instance_acc': test_instance_acc, 'class_acc': class_acc, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) global_epoch += 1 logger_loss.append([ train_loss, train_loss1, train_loss2, val_loss1, val_loss_wm1, val_loss2, val_loss_wm2 ]) logger_acc.append([ train_instance_acc, test_instance_acc1, instance_acc_wm, train_instance_acc2, test_instance_acc2, instance_acc_wm2 ]) time_end = datetime.datetime.now() time_span_str = str((time_end - time_start).seconds) log_string('Epoch time : %s S' % (time_span_str)) logger_loss.close() logger_loss.plot() savefig(os.path.join(log_dir, 'log_loss_v3.eps')) logger_acc.close() logger_acc.plot() savefig(os.path.join(log_dir, 'log_acc_v3.eps')) log_string('best_epoch %s' % str(best_epoch)) logger.info('End of training...')
def main(): args = parser.parse_args() torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) np.random.seed(args.seed) if args.remark != None: args.remark = args.remark else: args.remark = args.dataset + "-" + args.task + "-" + args.norm if args.dataset =="shapenet": args.num_class=16 else: args.num_class=40 def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) if args.task == 'baseline': experiment_dir_root = Path('/data-x/g12/zhangjie/3dIP/baseline') else: experiment_dir_root = Path('/data-x/g12/zhangjie/3dIP/ours') experiment_dir_root.mkdir(exist_ok=True) experiment_dir = experiment_dir_root.joinpath('pruning') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath(args.remark) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = experiment_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG_curve''' title = '' logger_loss = Logger(os.path.join(log_dir, 'log_loss_v3.txt'), title=title) logger_loss.set_names([ 'Valid Pub-Clean loss', 'Valid Pub-Trigger Loss', 'Valid Pri-Clean Loss', 'Valid Pri-Trigger Loss', ]) logger_acc = Logger(os.path.join(log_dir, 'log_acc_v3.txt'), title=title) logger_acc.set_names([ 'Model for Releasing', 'Model for Verification', 'Trigger', 'Signature']) '''LOG''' #创建log文件 logger = logging.getLogger("Model") #log的名字 logger.setLevel(logging.INFO) formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) #log的最低等级 file_handler.setFormatter(formatter) logger.addHandler(file_handler) #log文件名 log_string('PARAMETER ...') log_string(args) '''DATA LOADING''' log_string('Load pruning test dataset ...') if args.dataset == "shapenet": testDataLoader = getData.get_dataLoader(train=False, Shapenet=True, batchsize=args.batch_size) triggerDataLoader = getData2.get_dataLoader(Shapenet=True, T1=args.T1, batchsize=args.batch_size) else: testDataLoader = getData.get_dataLoader(train=False, Shapenet=False, batchsize=args.batch_size) triggerDataLoader = getData2.get_dataLoader(Shapenet=False,T1=args.T1, batchsize=args.batch_size) log_string('Load finished ...') '''MODEL LOADING''' num_class = args.num_class MODEL = importlib.import_module(args.model) shutil.copy('./models/%s.py' % args.model, str(experiment_dir)) shutil.copy('./models/pointnet_util.py', str(experiment_dir)) shutil.copy('prun3.py', str(experiment_dir)) classifier = MODEL.get_model(num_class, channel=3).cuda() # pprint(classifier) sd = experiment_dir_root.joinpath('classification') sd.mkdir(exist_ok=True) sd = sd.joinpath(str(args.remark)) sd.mkdir(exist_ok=True) sd = sd.joinpath('checkpoints/best_model.pth') log_string('pre-trained model chk pth: %s'%sd) checkpoint = torch.load(sd) model_dict = checkpoint['model_state_dict'] print('Total : {}'.format(len(model_dict))) print("best epoch", checkpoint['epoch']) classifier.load_state_dict(model_dict) classifier.cuda() p_num = get_parameter_number(classifier) log_string('Original trainable parameter: %s'%p_num) '''TESTING ORIGINAL''' logger.info('Test original model...') with torch.no_grad(): _, test_instance_acc1, class_acc1, _, _ = test(classifier, testDataLoader, num_class=args.num_class, ind=0) _, instance_acc_wm, class_acc_wm ,_, _ = test(classifier, triggerDataLoader, num_class=args.num_class, ind=0) _, test_instance_acc2, class_acc2, signloss2, signacc2 = test(classifier, testDataLoader, num_class=args.num_class,ind=1) _, instance_acc_wm2, class_acc_wm2, _, _ = test(classifier, triggerDataLoader,num_class=args.num_class, ind=1) log_string('Test Clean Public Accuracy: %f, Class Public Accuracy: %f' % (test_instance_acc1, class_acc1)) log_string('Test Clean Private Accuracy: %f, Class Private Accuracy: %f' % (test_instance_acc2, class_acc2)) log_string('Test Private Sign Accuracy: %f' % (signacc2)) log_string('Test Trigger Public Accuracy: %f, Trigger Class Public Accuracy: %f' % (instance_acc_wm, class_acc_wm)) log_string('Test Trigger Private Accuracy: %f, Trigger Class Private Accuracy: %f' % (instance_acc_wm2, class_acc_wm2)) '''PRUNING''' logger.info('Start testing of pruning...') for perc in [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]: time_start = datetime.datetime.now() classifier.load_state_dict(model_dict) p_num = get_parameter_number(classifier) log_string('Original trainable parameter: %s' % p_num) '''Testing pruning model''' logger.info('Testing pruning model--%d%%'%perc) pruning_net(classifier, perc) classifier.cuda() p_num = get_parameter_number(classifier) log_string('Pruning %02d%% -- trainable parameter: %s' % (perc, p_num)) with torch.no_grad(): val_loss1, test_instance_acc1, class_acc1, _, _ = test(classifier, testDataLoader, num_class=args.num_class, ind=0) val_loss_wm1, instance_acc_wm, class_acc_wm, _, _ = test(classifier, triggerDataLoader, num_class=args.num_class, ind=0) val_loss2, test_instance_acc2, class_acc2, signloss2, signacc2 = test(classifier, testDataLoader, num_class=args.num_class, ind=1) val_loss_wm2, instance_acc_wm2, class_acc_wm2, _, _ = test(classifier, triggerDataLoader, num_class=args.num_class, ind=1) log_string('Pruning %d%% Test Clean Public Accuracy: %f, Class Public Accuracy: %f' % (perc,test_instance_acc1, class_acc1)) log_string('Pruning %d%% Test Clean Private Accuracy: %f, Class Private Accuracy: %f' % (perc,test_instance_acc2, class_acc2)) log_string('Pruning %d%% Test Private Sign Accuracy: %f' % (perc,signacc2)) log_string('Pruning %d%% Test Trigger Public Accuracy: %f, Trigger Class Public Accuracy: %f' % (perc,instance_acc_wm, class_acc_wm)) log_string('Pruning %d%% Test Trigger Private Accuracy: %f, Trigger Class Private Accuracy: %f' % (perc,instance_acc_wm2, class_acc_wm2)) logger_loss.append([ val_loss1, val_loss_wm1, val_loss2,val_loss_wm2]) logger_acc.append([test_instance_acc1 *100 , test_instance_acc2 *100, instance_acc_wm*100, signacc2*100]) time_end = datetime.datetime.now() time_span_str = str((time_end - time_start).seconds) log_string('Epoch time : %s S' % (time_span_str)) logger_loss.close() logger_loss.plot_prun() savefig(os.path.join(log_dir, 'log_loss.eps')) logger_acc.close() logger_acc.plot_prun() acc_name = args.remark + '_prun.eps' savefig(os.path.join(log_dir, acc_name)) logger.info('End of pruning...')
def main(): args = parser.parse_args() torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) np.random.seed(args.seed) if args.remark != None: args.remark = args.remark else: args.remark = args.dataset + "-" + args.task + "-" + args.norm if args.dataset == "shapenet": args.num_class = 16 else: args.num_class = 40 def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) # experiment_dir = Path('./exp/v2/') experiment_dir = Path('/data-x/g12/zhangjie/3dIP/exp3.0/v2') # experiment_dir = Path('/data-x/g10/zhangjie/3D/exp/v2') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath('classification') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath(args.remark + "_" + timestr) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = experiment_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG_curve''' title = args.dataset + "-" + args.task + "-" + args.norm logger_loss = Logger(os.path.join(log_dir, 'log_loss.txt'), title=title) logger_loss.set_names([ 'Train AVE Loss', 'Train Public Loss', 'Train Private Loss', 'Valid AVE Loss', 'Valid Public Loss', 'Valid Private Loss' ]) logger_acc = Logger(os.path.join(log_dir, 'log_acc.txt'), title=title) logger_acc.set_names([ 'Train Public Acc.', 'Train Private Acc.', 'Test Public Acc.', 'Test Private Acc.' ]) '''LOG''' #创建log文件 logger = logging.getLogger("Model") #log的名字 # print("FFFFFFFF",logger) #<Logger Model (WARNING)> logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) #log的最低等级 file_handler.setFormatter(formatter) logger.addHandler(file_handler) #log文件名 log_string('PARAMETER ...') # print("FFFFFFF",logger.info) #<bound method Logger.info of <Logger Model (INFO)>> log_string(args) '''DATA LOADING''' log_string('Load dataset ...') if args.dataset == "shapenet": trainDataLoader = getData.get_dataLoader(train=True, Shapenet=True, batchsize=args.batch_size) testDataLoader = getData.get_dataLoader(train=False, Shapenet=True, batchsize=args.batch_size) else: trainDataLoader = getData.get_dataLoader(train=True, Shapenet=False, batchsize=args.batch_size) testDataLoader = getData.get_dataLoader(train=False, Shapenet=False, batchsize=args.batch_size) log_string('Finished ...') log_string('Load model ...') '''MODEL LOADING''' num_class = args.num_class MODEL = importlib.import_module(args.model) # 当在写代码时,我们希望能够根据传入的选项设置,如args.model来确定要导入使用的是哪个model.py文件,而不是一股脑地导入, 这种时候就需要用上python的动态导入模块 # 复制model文件到exp——dir shutil.copy('./models/%s.py' % args.model, str(experiment_dir)) shutil.copy('./models/pointnet_util.py', str(experiment_dir)) shutil.copy('train_2_cls.py', str(experiment_dir)) shutil.copy('./data/getData.py', str(experiment_dir)) shutil.copytree('./models/layers', str(experiment_dir) + "/layers") classifier = MODEL.get_model(num_class, channel=3).cuda() criterion = MODEL.get_loss().cuda() pprint(classifier) try: checkpoint = torch.load( str(experiment_dir) + '/checkpoints/best_model.pth') start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) log_string('Use pretrain model') except: log_string('No existing model, starting training from scratch...') start_epoch = 0 if args.optimizer == 'Adam': optimizer = torch.optim.Adam(classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate) else: optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7) global_epoch = 0 global_step = 0 best_instance_acc = 0.0 best_class_acc = 0.0 mean_correct = [] mean_correct2 = [] mean_loss = [] mean_loss1 = [] mean_loss2 = [] '''TRANING''' logger.info('Start training...') for epoch in range(start_epoch, args.epoch): time_start = datetime.datetime.now() log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) scheduler.step() for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): points, target = data points = points.data.numpy() points = provider.random_point_dropout( points) #provider是自己写的一个对点云操作的函数,随机dropout,置为第一个点的值 points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3]) #点的放缩 points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3]) #点的偏移 points = torch.Tensor(points) points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() optimizer.zero_grad() classifier = classifier.train() for m in classifier.modules(): if isinstance(m, SignLoss): m.reset() loss1 = torch.tensor(0.).cuda() loss2 = torch.tensor(0.).cuda() sign_loss = torch.tensor(0.).cuda() for ind in range(2): if ind == 0: pred, trans_feat = classifier(points, ind=ind) loss1 = criterion(pred, target.long(), trans_feat) mean_loss1.append(loss1.item() / float(points.size()[0])) pred_choice = pred.data.max(1)[1] correct = pred_choice.eq(target.long().data).cpu().sum() mean_correct.append(correct.item() / float(points.size()[0])) else: pred2, trans_feat2 = classifier(points, ind=ind) loss2 = criterion(pred2, target.long(), trans_feat2) mean_loss2.append(loss2.item() / float(points.size()[0])) pred_choice2 = pred2.data.max(1)[1] correct2 = pred_choice2.eq(target.long().data).cpu().sum() mean_correct2.append(correct2.item() / float(points.size()[0])) for m in classifier.modules(): if isinstance(m, SignLoss): sign_loss += m.loss loss = args.beta * loss1 + loss2 + sign_loss mean_loss.append(loss.item() / float(points.size()[0])) # loss = loss2 loss.backward() optimizer.step() global_step += 1 train_instance_acc = np.mean(mean_correct) train_instance_acc2 = np.mean(mean_correct2) train_instance_acc_ave = (train_instance_acc + train_instance_acc2) / 2 train_loss = np.mean(mean_loss) / 2 train_loss1 = np.mean(mean_loss1) train_loss2 = np.mean(mean_loss2) log_string('Train Instance Public Accuracy: %f' % train_instance_acc) log_string('Train Instance Private Accuracy: %f' % train_instance_acc2) sign_acc = torch.tensor(0.).cuda() count = 0 for m in classifier.modules(): if isinstance(m, SignLoss): sign_acc += m.acc count += 1 if count != 0: sign_acc /= count log_string('Sign Accuracy: %f' % sign_acc) with torch.no_grad(): for ind in range(2): if ind == 0: val_loss1, test_instance_acc1, class_acc1 = test( classifier, testDataLoader, num_class=args.num_class, ind=0) else: val_loss2, test_instance_acc2, class_acc2 = test( classifier, testDataLoader, num_class=args.num_class, ind=1) log_string( 'Test Instance Public Accuracy: %f, Class Public Accuracy: %f' % (test_instance_acc1, class_acc1)) log_string( 'Test Instance Private Accuracy: %f, Class Private Accuracy: %f' % (test_instance_acc2, class_acc2)) val_loss = (val_loss1 + val_loss2) / 2 test_instance_acc = (test_instance_acc1 + test_instance_acc2) / 2 class_acc = (class_acc1 + class_acc2) / 2 if (test_instance_acc >= best_instance_acc): best_instance_acc = test_instance_acc best_epoch = epoch + 1 if (class_acc >= best_class_acc): best_class_acc = class_acc log_string( 'Test Instance Average Accuracy: %f, Class Average Accuracy: %f' % (test_instance_acc, class_acc)) log_string( 'Best Instance Average Accuracy: %f, Class Average Accuracy: %f' % (best_instance_acc, best_class_acc)) if (test_instance_acc >= best_instance_acc): logger.info('Save model...') savepath = str(checkpoints_dir) + '/best_model.pth' log_string('Saving at %s' % savepath) log_string('best_epoch %s' % str(best_epoch)) state = { 'epoch': best_epoch, 'instance_acc': test_instance_acc, 'class_acc': class_acc, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) global_epoch += 1 logger_loss.append([ train_loss, train_loss1, train_loss2, val_loss, val_loss1, val_loss2 ]) logger_acc.append([ train_instance_acc, train_instance_acc2, test_instance_acc1, test_instance_acc2 ]) time_end = datetime.datetime.now() time_span_str = str((time_end - time_start).seconds) log_string('Epoch time : %s S' % (time_span_str)) logger_loss.close() logger_loss.plot() savefig(os.path.join(log_dir, 'log_loss.eps')) logger_acc.close() logger_acc.plot() savefig(os.path.join(log_dir, 'log_acc.eps')) log_string('best_epoch %s' % str(best_epoch)) logger.info('End of training...')