p.add_argument('--checkpoint_path',
default=None,
help='Checkpoint to trained model.')
p.add_argument('--conv_encoder',
action='store_true',
default=False,
help='Use convolutional encoder process')
opt = p.parse_args()
assert opt.dataset == 'celeba_32x32'
if opt.conv_encoder: gmode = 'conv_cnp'
else: gmode = 'cnp'
img_dataset = dataio.CelebA(split='train', downsampled=True)
coord_dataset = dataio.Implicit2DWrapper(
img_dataset,
sidelength=(32, 32),
train_sparsity_range=opt.train_sparsity_range,
generalization_mode=gmode)
image_resolution = (32, 32)
dataloader = DataLoader(coord_dataset,
shuffle=True,
batch_size=opt.batch_size,
pin_memory=True,
num_workers=0)
if opt.conv_encoder:
model = meta_modules.ConvolutionalNeuralProcessImplicit2DHypernet(
default=200,
help='Amount of subsampled pixels input into the set encoder')
p.add_argument('--partial_conv',
action='store_true',
default=False,
help='Use a partial convolution encoder')
opt = p.parse_args()
if opt.experiment_name is None:
opt.experiment_name = opt.checkpoint_path.split('/')[-3] + '_TEST'
else:
opt.experiment_name = opt.checkpoint_path.split(
'/')[-3] + '_' + opt.experiment_name
assert opt.dataset == 'celeba_32x32'
img_dataset_test = dataio.CelebA(split='test', downsampled=True)
coord_dataset_test = dataio.Implicit2DWrapper(
img_dataset_test,
sidelength=(32, 32),
test_sparsity=200,
generalization_mode='conv_cnp_test')
image_resolution = (32, 32)
img_dataset_train = dataio.CelebA(split='train', downsampled=True)
coord_dataset_train = dataio.Implicit2DWrapper(
img_dataset_train,
sidelength=(32, 32),
test_sparsity=200,
generalization_mode='conv_cnp_test')
# Define the model.