first pretrain/load a mdoel, then use to craft poisoned dataset under the priciple of gradient alignment, then used the poisoned dataset to train a new model and use the poisoned new model to craft better poison dataset
(self, init_model, schedule=None)
| 476 | return poison_trainset, patch_source_testset, full_patch_testset |
| 477 | |
| 478 | def train(self, init_model, schedule=None): |
| 479 | """first pretrain/load a mdoel, then use to craft poisoned dataset under the priciple of gradient alignment, then used the poisoned dataset to train a new model and use the poisoned new model to craft better poison dataset""" |
| 480 | if schedule is None and self.global_schedule is None: |
| 481 | raise AttributeError("Training schedule is None, please check your schedule setting.") |
| 482 | elif schedule is not None and self.global_schedule is None: |
| 483 | self.current_schedule = deepcopy(schedule) |
| 484 | elif schedule is None and self.global_schedule is not None: |
| 485 | self.current_schedule = deepcopy(self.global_schedule) |
| 486 | elif schedule is not None and self.global_schedule is not None: |
| 487 | self.current_schedule = deepcopy(schedule) |
| 488 | |
| 489 | if 'pretrain' in self.current_schedule and os.path.exists(self.current_schedule['pretrain']): |
| 490 | self.model.load_state_dict(torch.load(self.current_schedule['pretrain']), strict=False) |
| 491 | |
| 492 | # Select Device |
| 493 | # Use GPU |
| 494 | if 'device' in self.current_schedule and self.current_schedule['device'] == 'GPU': |
| 495 | if 'CUDA_VISIBLE_DEVICES' in self.current_schedule: |
| 496 | os.environ['CUDA_VISIBLE_DEVICES'] = self.current_schedule['CUDA_VISIBLE_DEVICES'] |
| 497 | |
| 498 | assert torch.cuda.device_count() > 0, 'This machine has no cuda devices!' |
| 499 | assert self.current_schedule['GPU_num'] >0, 'GPU_num should be a positive integer' |
| 500 | print(f"This machine has {torch.cuda.device_count()} cuda devices, and use {self.current_schedule['GPU_num']} of them to train.") |
| 501 | |
| 502 | if self.current_schedule['GPU_num'] == 1: |
| 503 | device = torch.device("cuda:0") |
| 504 | else: |
| 505 | gpus = list(range(self.current_schedule['GPU_num'])) |
| 506 | self.model = nn.DataParallel(self.model.cuda(), device_ids=gpus, output_device=gpus[0]) |
| 507 | # TODO: DDP training |
| 508 | pass |
| 509 | # Use CPU |
| 510 | else: |
| 511 | device = torch.device("cpu") |
| 512 | work_dir = osp.join(self.current_schedule['save_dir'], self.current_schedule['experiment_name'] + '_' + time.strftime("%Y-%m-%d_%H:%M:%S", time.localtime())) |
| 513 | os.makedirs(work_dir, exist_ok=True) |
| 514 | log = Log(osp.join(work_dir, 'log.txt')) |
| 515 | |
| 516 | self.model = self.model.to(device) |
| 517 | self.model.train() |
| 518 | |
| 519 | h = self.train_dataset[0][0].shape[1] |
| 520 | augment = RandomTransform(source_size=h, target_size=h, shift=h//4) |
| 521 | |
| 522 | |
| 523 | if self.current_schedule['benign_training'] is True: |
| 524 | self.current_schedule['milestones']=self.current_schedule['schedule'] |
| 525 | self._train_model(self.model, log, self.train_dataset, self.test_dataset, None, None ,augment, device, self.current_schedule) |
| 526 | elif self.current_schedule['benign_training'] is False: |
| 527 | _, _, _, patch_source_testset, patch_testset = prepare_dataset(self.source_num, self.train_dataset, self.test_dataset, self.y_target, self.y_source, self.patch, self.random_patch) |
| 528 | log("******pretraining*********\n") |
| 529 | if ('pretrain' not in self.current_schedule) or ('pretrain' in self.current_schedule and not os.path.exists(self.current_schedule['pretrain'])): |
| 530 | self._train_model(model=self.model, |
| 531 | log=log, |
| 532 | trainset=self.train_dataset, |
| 533 | testset=self.test_dataset, |
| 534 | poison_sourceset=patch_source_testset, |
| 535 | poison_testset=patch_testset, |
no test coverage detected