def __init__(self, pgd_alpha: float = 1.0, pgd_epsilon: float = 8.0, pgd_iteration: int = 8, stop_conf: float = 0.9, magnet: bool = False, randomized_smooth: bool = False, curvature: bool = False, **kwargs): super().__init__(**kwargs) self.param_list['pgd'] = ['pgd_alpha', 'pgd_epsilon', 'pgd_iteration'] self.pgd_alpha: float = pgd_alpha self.pgd_epsilon: float = pgd_epsilon self.pgd_iteration: int = pgd_iteration self.pgd = PGD_Optimizer(alpha=self.pgd_alpha / 255, epsilon=self.pgd_epsilon / 255, iteration=self.pgd_iteration) self.stop_conf: float = stop_conf if magnet: self.magnet: MagNet = MagNet(dataset=self.dataset, pretrain=True) self.randomized_smooth: bool = randomized_smooth if curvature: self.curvature: Curvature = Curvature(model=self.model)
def __init__(self, noise_dim: int, dim: int, data_shape: list[int] = [3, 32, 32], generator_iters: int = 1000, critic_iter: int = 5): self.noise_dim = noise_dim self.G: Generator = Generator(noise_dim, dim, data_shape) self.D: Discriminator = Discriminator(dim, data_shape) if env['num_gpus']: self.G.cuda() self.D.cuda() # the parameter in the original paper self.d_optimizer = optim.RMSprop(self.D.parameters(), lr=5e-5) self.g_optimizer = optim.RMSprop(self.G.parameters(), lr=5e-5) self.generator_iters = generator_iters # larger: 1000 self.critic_iter = critic_iter self.mse_loss = torch.nn.MSELoss() self.gan_pgd: PGD_Optimizer = PGD_Optimizer(epsilon=1.0, iteration=500, output=0)