])),
batch_size=args.val_batchsize,
shuffle=False,
num_workers=args.nworkers,
)
elif args.data == 'celebahq':
im_dim = 3
init_layer = layers.LogitTransform(0.05)
if args.imagesize != 256:
logger.info('Changing image size to 256.')
args.imagesize = 256
train_loader = torch.utils.data.DataLoader(datasets.CelebAHQ(
train=True,
transform=transforms.Compose([
transforms.ToPILImage(),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
reduce_bits,
lambda x: add_noise(x, nvals=2**args.nbits),
])),
batch_size=args.batchsize,
shuffle=True,
num_workers=args.nworkers)
test_loader = torch.utils.data.DataLoader(datasets.CelebAHQ(
train=False,
transform=transforms.Compose([
reduce_bits,
lambda x: add_noise(x, nvals=2**args.nbits),
])),
batch_size=args.val_batchsize,
shuffle=False,
batch_size=args.val_batchsize,
shuffle=False,
num_workers=args.nworkers,
)
elif args.data == 'celebahq':
im_dim = 3
init_layer = layers.LogitTransform(0.05)
if args.imagesize != 256:
logger.info('Changing image size to 256.')
args.imagesize = 256
train_loader = torch.utils.data.DataLoader(
datasets.CelebAHQ(
train=True, transform=transforms.Compose([
transforms.ToPILImage(),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
reduce_bits,
lambda x: add_noise(x, nvals=2**args.nbits),
])
), batch_size=args.batchsize, shuffle=True, num_workers=args.nworkers
)
test_loader = torch.utils.data.DataLoader(
datasets.CelebAHQ(
train=False, transform=transforms.Compose([
reduce_bits,
lambda x: add_noise(x, nvals=2**args.nbits),
])
), batch_size=args.val_batchsize, shuffle=False, num_workers=args.nworkers
)
elif args.data == 'celeba_5bit':
im_dim = 3
checkpoint = torch.load(model_name)
args = checkpoint['args']
args.val_batchsize = 6
args.batchsize = 6
# define data
if args.data == 'celeba64_5bit':
im_dim = 3
init_layer = layers.LogitTransform(0.05)
if args.imagesize != 64:
args.imagesize = 64
# no dequantization for attack!
train_loader = torch.utils.data.DataLoader(datasets.CelebAHQ(
train=True,
transform=transforms.Compose([
transforms.ToPILImage(),
transforms.Resize(args.imagesize),
transforms.ToTensor(), reduce_bits
])),
batch_size=args.batchsize,
shuffle=False,
num_workers=args.nworkers)
# no dequantization for attack!
test_loader = torch.utils.data.DataLoader(datasets.CelebAHQ(
train=False,
transform=transforms.Compose(
[transforms.ToPILImage(),
transforms.ToTensor(), reduce_bits])),
batch_size=args.val_batchsize,
shuffle=False,
num_workers=args.nworkers)
