def add_argument(parser: argparse.ArgumentParser, dataset_name: str = None, dataset: str | Dataset = None, config: Config = config, class_dict: dict[str, type[Dataset]] = {} ) -> argparse._ArgumentGroup: r""" | Add dataset arguments to argument parser. | For specific arguments implementation, see :meth:`Dataset.add_argument()`. Args: parser (argparse.ArgumentParser): The parser to add arguments. dataset_name (str): The dataset name. dataset (str | Dataset): Dataset instance or dataset name (as the alias of `dataset_name`). config (Config): The default parameter config, which contains the default dataset name if not provided. class_dict (dict[str, type[Dataset]]): Map from dataset name to dataset class. Defaults to ``{}``. """ dataset_name = get_name( name=dataset_name, module=dataset, arg_list=['-d', '--dataset']) dataset_name = dataset_name if dataset_name is not None \ else config.full_config['dataset']['default_dataset'] group = parser.add_argument_group( '{yellow}dataset{reset}'.format(**ansi), description=dataset_name) try: DatasetType = class_dict[dataset_name] except KeyError: print(f'{dataset_name} not in \n{list(class_dict.keys())}') raise return DatasetType.add_argument(group)
def add_argument(parser: argparse.ArgumentParser, defense_name: None | str = None, defense: None | str | Defense = None, class_dict: dict[str, type[Defense]] = {}): r""" | Add defense arguments to argument parser. | For specific arguments implementation, see :meth:`Defense.add_argument()`. Args: parser (argparse.ArgumentParser): The parser to add arguments. defense_name (str): The defense name. defense (str | Defense): The defense instance or defense name (as the alias of `defense_name`). class_dict (dict[str, type[Defense]]): Map from defense name to defense class. Defaults to ``{}``. Returns: argparse._ArgumentGroup: The argument group. """ defense_name = get_name(name=defense_name, module=defense, arg_list=['--defense']) group = parser.add_argument_group('{yellow}defense{reset}'.format(**ansi), description=defense_name) try: DefenseType = class_dict[defense_name] except KeyError: if defense_name is None: print(f'{ansi["red"]}you need to first claim the defense name ' f'using "--defense".{ansi["reset"]}') print(f'{defense_name} not in \n{list(class_dict.keys())}') raise return DefenseType.add_argument(group)
def add_argument(parser: argparse.ArgumentParser, attack_name: str = None, attack: str | Attack = None, class_dict: dict[str, type[Attack]] = {}): r""" | Add attack arguments to argument parser. | For specific arguments implementation, see :meth:`Attack.add_argument()`. Args: parser (argparse.ArgumentParser): The parser to add arguments. attack_name (str): The attack name. attack (str | Attack): The attack instance or attack name (as the alias of `attack_name`). class_dict (dict[str, type[Attack]]): Map from attack name to attack class. Defaults to ``{}``. Returns: argparse._ArgumentGroup: The argument group. """ attack_name = get_name(name=attack_name, module=attack, arg_list=['--attack']) group = parser.add_argument_group('{yellow}attack{reset}'.format(**ansi), description=attack_name) try: AttackType = class_dict[attack_name] except KeyError: if attack_name is None: print(f'{ansi["red"]}you need to first claim the attack name ' f'using "--attack".{ansi["reset"]}') print(f'{attack_name} not in \n{list(class_dict.keys())}') raise return AttackType.add_argument(group)
def create(dataset_name: str = None, dataset: str = None, config: Config = config, class_dict: dict[str, type[Dataset]] = {}, **kwargs) -> Dataset: r""" | Create a dataset instance. | For arguments not included in :attr:`kwargs`, use the default values in :attr:`config`. | The default value of :attr:`folder_path` is ``'{data_dir}/{data_type}/{name}'``. | For dataset implementation, see :class:`Dataset`. Args: dataset_name (str): The dataset name. dataset (str): The alias of `dataset_name`. config (Config): The default parameter config. class_dict (dict[str, type[Dataset]]): Map from dataset name to dataset class. Defaults to ``{}``. **kwargs: Keyword arguments passed to dataset init method. Returns: Dataset: Dataset instance. """ dataset_name = get_name( name=dataset_name, module=dataset, arg_list=['-d', '--dataset']) dataset_name = dataset_name if dataset_name is not None \ else config.full_config['dataset']['default_dataset'] result = config.get_config(dataset_name=dataset_name)[ 'dataset'].update(kwargs) try: DatasetType = class_dict[dataset_name] except KeyError: print(f'{dataset_name} not in \n{list(class_dict.keys())}') raise if 'folder_path' not in result.keys(): result['folder_path'] = os.path.join(result['data_dir'], DatasetType.data_type, DatasetType.name) return DatasetType(**result)
def create(cmd_config_path: str = None, dataset_name: str = None, dataset: str = None, seed: int = None, data_seed: int = None, cudnn_benchmark: bool = None, config: Config = config, cache_threshold: float = None, verbose: int = 0, color: bool = None, device: str | int | torch.device = None, tqdm: bool = None, **kwargs) -> Env: r""" | Load :attr:`env` values from config and command line. Args: dataset_name (str): The dataset name. dataset (str | trojanzoo.datasets.Dataset): Dataset instance (required for :attr:`model_ema`) or dataset name (as the alias of `dataset_name`). model (trojanzoo.models.Model): Model instance. config (Config): The default parameter config. **kwargs: The keyword arguments in keys of ``['optim_args', 'train_args', 'writer_args']``. Returns: Env: The :attr:`env` instance. """ other_kwargs = {'data_seed': data_seed, 'cache_threshold': cache_threshold, 'verbose': verbose, 'color': color, 'device': device, 'tqdm': tqdm} config.cmd_config_path = cmd_config_path dataset_name = get_name( name=dataset_name, module=dataset, arg_list=['-d', '--dataset']) dataset_name = dataset_name if dataset_name is not None \ else config.full_config['dataset']['default_dataset'] result = config.get_config(dataset_name=dataset_name)[ 'env'].update(other_kwargs) env.clear().update(**result) ansi.switch(env['color']) if seed is None and 'seed' in env.keys(): seed = env['seed'] random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) num_gpus: int = torch.cuda.device_count() device: str | int | torch.device = result['device'] if device is None: device = 'auto' match device: case torch.device(): pass case 'auto': device = torch.device('cuda' if num_gpus else 'cpu') case 'gpu': device = torch.device('cuda') case _: device = torch.device(device) if device.type == 'cpu': num_gpus = 0 if device.index is not None and torch.cuda.is_available(): num_gpus = 1 if cudnn_benchmark is None and 'cudnn_benchmark' in env.keys(): cudnn_benchmark = env['cudnn_benchmark'] if cudnn_benchmark: torch.backends.cudnn.benchmark = cudnn_benchmark env.update(seed=seed, device=device, cudnn_benchmark=cudnn_benchmark, num_gpus=num_gpus) env['world_size'] = 1 # TODO return env
def create(defense_name: None | str = None, defense: None | str | Defense = None, folder_path: None | str = None, dataset_name: str = None, dataset: None | str | Dataset = None, model_name: str = None, model: None | str | Model = None, config: Config = config, class_dict: dict[str, type[Defense]] = {}, **kwargs): r""" | Create a defense instance. | For arguments not included in :attr:`kwargs`, use the default values in :attr:`config`. | The default value of :attr:`folder_path` is ``'{defense_dir}/{dataset.data_type}/{dataset.name}/{model.name}/{defense.name}'``. | For defense implementation, see :class:`Defense`. Args: defense_name (str): The defense name. defense (str | Defense): The defense instance or defense name (as the alias of `defense_name`). dataset_name (str): The dataset name. dataset (str | trojanzoo.datasets.Dataset): Dataset instance or dataset name (as the alias of `dataset_name`). model_name (str): The model name. model (str | Model): The model instance or model name (as the alias of `model_name`). config (Config): The default parameter config. class_dict (dict[str, type[Defense]]): Map from defense name to defense class. Defaults to ``{}``. **kwargs: The keyword arguments passed to defense init method. Returns: Defense: The defense instance. """ dataset_name = get_name(name=dataset_name, module=dataset, arg_list=['-d', '--dataset']) model_name = get_name(name=model_name, module=model, arg_list=['-m', '--model']) defense_name = get_name(name=defense_name, module=defense, arg_list=['--defense']) if dataset_name is None: dataset_name = config.full_config['dataset']['default_dataset'] general_config = config.get_config(dataset_name=dataset_name)['defense'] specific_config = config.get_config( dataset_name=dataset_name)[defense_name] result = general_config.update(specific_config).update( kwargs) # TODO: linting issues try: DefenseType = class_dict[defense_name] except KeyError: print(f'{defense_name} not in \n{list(class_dict.keys())}') raise if 'folder_path' not in result.keys(): folder_path = result['defense_dir'] if isinstance(dataset, Dataset): folder_path = os.path.join(folder_path, dataset.data_type, dataset.name) if model_name is not None: folder_path = os.path.join(folder_path, model_name) folder_path = os.path.join(folder_path, DefenseType.name) result['folder_path'] = folder_path return DefenseType(name=defense_name, dataset=dataset, model=model, **result)
def create(attack_name: str = None, attack: str | Attack = None, dataset_name: str = None, dataset: str | Dataset = None, model_name: str = None, model: str | Model = None, config: Config = config, class_dict: dict[str, type[Attack]] = {}, **kwargs) -> Attack: r""" | Create an attack instance. | For arguments not included in :attr:`kwargs`, use the default values in :attr:`config`. | The default value of :attr:`folder_path` is ``'{attack_dir}/{dataset.data_type}/{dataset.name}/{model.name}/{attack.name}'``. | For attack implementation, see :class:`Attack`. Args: attack_name (str): The attack name. attack (str | Attack): The attack instance or attack name (as the alias of `attack_name`). dataset_name (str): The dataset name. dataset (str | Dataset): Dataset instance or dataset name (as the alias of `dataset_name`). model_name (str): The model name. model (str | Model): The model instance or model name (as the alias of `model_name`). config (Config): The default parameter config. class_dict (dict[str, type[Attack]]): Map from attack name to attack class. Defaults to ``{}``. **kwargs: The keyword arguments passed to attack init method. Returns: Attack: The attack instance. """ dataset_name = get_name(name=dataset_name, module=dataset, arg_list=['-d', '--dataset']) model_name = get_name(name=model_name, module=model, arg_list=['-m', '--model']) attack_name = get_name(name=attack_name, module=attack, arg_list=['--attack']) if dataset_name is None: dataset_name = config.full_config['dataset']['default_dataset'] general_config = config.get_config(dataset_name=dataset_name)['attack'] specific_config = config.get_config(dataset_name=dataset_name)[attack_name] result = general_config.update(specific_config).update(kwargs) try: AttackType = class_dict[attack_name] except KeyError: print(f'{attack_name} not in \n{list(class_dict.keys())}') raise if 'folder_path' not in result.keys(): folder_path = result['attack_dir'] if isinstance(dataset, Dataset): folder_path = os.path.join(folder_path, dataset.data_type, dataset.name) if model_name is not None: folder_path = os.path.join(folder_path, model_name) folder_path = os.path.join(folder_path, AttackType.name) result['folder_path'] = folder_path return AttackType(name=attack_name, dataset=dataset, model=model, **result)
def create(dataset_name: None | str = None, dataset: None | str | Dataset = None, model: None | Model = None, model_ema: None | bool = False, pre_conditioner: None | str = None, tensorboard: None | bool = None, ClassType: type[Trainer] = Trainer, config: Config = config, **kwargs): r""" | Create a trainer instance. | For arguments not included in :attr:`kwargs`, use the default values in :attr:`config`. | For trainer implementation, see :class:`Trainer`. Args: dataset_name (str): The dataset name. dataset (str | trojanzoo.datasets.Dataset): Dataset instance (required for :attr:`model_ema`) or dataset name (as the alias of `dataset_name`). model (trojanzoo.models.Model): Model instance. model_ema (bool): Whether to use :class:`~trojanzoo.utils.model.ExponentialMovingAverage`. Defaults to ``False``. pre_conditioner (str): Choose from * ``None`` * ``'kfac'``: :class:`~trojanzoo.utils.fim.KFAC` * ``'ekfac'``: :class:`~trojanzoo.utils.fim.EKFAC` Defaults to ``None``. tensorboard (bool): Whether to use :any:`torch.utils.tensorboard.writer.SummaryWriter`. Defaults to ``False``. ClassType (type[Trainer]): The trainer class type. Defaults to :class:`Trainer`. config (Config): The default parameter config. **kwargs: The keyword arguments in keys of ``['optim_args', 'train_args', 'writer_args']``. Returns: Trainer: The trainer instance. """ assert isinstance(model, Model) dataset_name = get_name(name=dataset_name, module=dataset, arg_list=['-d', '--dataset']) result = config.get_config(dataset_name=dataset_name )['trainer'].update(kwargs) optim_keys = model.define_optimizer.__code__.co_varnames train_keys = model._train.__code__.co_varnames optim_args: dict[str, Any] = {} train_args: dict[str, Any] = {} for key, value in result.items(): if key in optim_keys: optim_args[key] = value elif key in train_keys and key != 'verbose': train_args[key] = value train_args['epochs'] = result['epochs'] train_args['lr_warmup_epochs'] = result['lr_warmup_epochs'] optimizer, lr_scheduler = model.define_optimizer(**optim_args) module = model._model match optim_args['parameters']: case 'features': module = module.features case 'classifier': module = module.classifier # https://github.com/pytorch/vision/blob/main/references/classification/train.py model_ema_module = None if model_ema: # Decay adjustment that aims to keep the decay independent from other hyper-parameters originally proposed at: # https://github.com/facebookresearch/pycls/blob/f8cd9627/pycls/core/net.py#L123 # # total_ema_updates = (Dataset_size / n_GPUs) * epochs / (batch_size_per_gpu * EMA_steps) # We consider constant = Dataset_size for a given dataset/setup and ommit it. Thus: # adjust = 1 / total_ema_updates ~= n_GPUs * batch_size_per_gpu * EMA_steps / epochs adjust = env['world_size'] * dataset.batch_size * \ result['model_ema_steps'] / result['epochs'] alpha = 1.0 - result['model_ema_decay'] alpha = min(1.0, alpha * adjust) model_ema_module = ExponentialMovingAverage( model._model, decay=1.0 - alpha) match pre_conditioner: case 'kfac': kfac_optimizer = KFAC(module) case 'ekfac': kfac_optimizer = EKFAC(module) case _: kfac_optimizer = None writer = None writer_args: dict[str, Any] = {} if tensorboard: from torch.utils.tensorboard import SummaryWriter # log_dir, flush_secs, ... writer_keys = SummaryWriter.__init__.__code__.co_varnames for key, value in result.items(): if key in writer_keys: writer_args[key] = value writer = SummaryWriter(**writer_args) return ClassType(optim_args=optim_args, train_args=train_args, writer_args=writer_args, optimizer=optimizer, lr_scheduler=lr_scheduler, model_ema=model_ema_module, pre_conditioner=kfac_optimizer, writer=writer)