def main(run_experiment):
"""
Runs encrypted contextual bandits learning experiment on MNIST.
"""
# parse input arguments:
args = parse_args(os.environ.get("HOSTNAME", "localhost"))
if args.mnist_dir is not None:
process_mnist_files(args.mnist_dir, "/tmp/MNIST/processed")
if args.multiprocess:
launcher = MultiProcessLauncher(args.world_size, run_experiment, args)
launcher.start()
launcher.join()
launcher.terminate()
else:
run_experiment(args)
def preprocess_data(context_manager, data_dirname):
if mnist_dir is not None:
process_mnist_files(
mnist_dir, os.path.join(data_dirname, "MNIST", "processed")
)
download = False
else:
download = True
with context_manager:
if not evaluate:
mnist_train = datasets.MNIST(
data_dirname,
download=download,
train=True,
transform=transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,)),
]
),
)
mnist_test = datasets.MNIST(
data_dirname,
download=download,
train=False,
transform=transforms.Compose(
[transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
),
)
train_loader = (
torch.utils.data.DataLoader(
mnist_train, batch_size=batch_size, shuffle=True
)
if not evaluate
else None
)
test_loader = torch.utils.data.DataLoader(
mnist_test, batch_size=batch_size, shuffle=False
)
return train_loader, test_loader