Pytorch implementation of Neural Processes. This repo follows the best practices defined in Empirical Evaluation of Neural Process Objectives.
Simple example of training a neural process on functions or images.
import torch from neural_process import NeuralProcess, NeuralProcessImg from training import NeuralProcessTrainer # Define neural process for functions... neuralprocess = NeuralProcess(x_dim=1, y_dim=1, r_dim=10, z_dim=10, h_dim=10) # ...or for images neuralprocess = NeuralProcessImg(img_size=(3, 32, 32), r_dim=128, z_dim=128, h_dim=128) # Define optimizer and trainer optimizer = torch.optim.Adam(neuralprocess.parameters(), lr=3e-4) np_trainer = NeuralProcessTrainer(device, neuralprocess, optimizer, num_context_range=(3, 20), num_extra_target_range=(5, 10)) # Train on your data np_trainer.train(data_loader, epochs=30)
For a detailed tutorial on training and using neural processes on 1d functions, see
To train an image model, use
python main_experiment.py config.json. This will log information about training and save model weights.
For a detailed tutorial on how to load a trained model and how to use neural processes for inpainting, see the notebook
example-img. Trained models for MNIST and CelebA are also provided in the
Note, to train on CelebA you will have to download the data from here.