<jupyter_start><jupyter_text>A toy example of Real NVP (using the ML class)<jupyter_code>from __future__ import print_function
import torch
import torch.utils.data
from torch import nn, optim
from torchvision import datasets

from matplotlib import pyplot as plt

batch_size = 128
epochs = 5000
test_size = 500
seed = 1
torch.manual_seed(seed)

if torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"
from pixyz.distributions import Normal, InverseTransformedDistribution
from pixyz.flows import AffineCoupling, FlowList, BatchNorm1d
from pixyz.models import ML
from pixyz.utils import print_latex
def plot_samples(points, noise):
    X_LIMS = (-1.5, 2.5)
    Y_LIMS = (-2.5, 2.5)

    fig = plt.figure(figsize=(8, 4))
    ax = fig.add_subplot(121)
    ax.scatter(points[:, 0], points[:, 1], alpha=0.7, s=25, c="b")
    ax.set_xlim(*X_LIMS)
    ax.set_ylim(*Y_LIMS)
    ax.set_xlabel("p(x)")
    
    X_LIMS = (-3, 3)
    Y_LIMS = (-3, 3)
    
    ax = fig.add_subplot(122)
    ax.scatter(noise[:, 0], noise[:, 1], alpha=0.7, s=25, c="r")
    ax.set_xlim(*X_LIMS)
    ax.set_ylim(*Y_LIMS)
    ax.set_xlabel("p(z)")

    plt.show()
x_dim = 2
z_dim = x_dim
# prior
prior = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.),
               var=["z"], features_shape=[z_dim], name="prior").to(device)
class ScaleTranslateNet(nn.Module):
    def __init__(self, in_features, hidden_features):
        super().__init__()
        self.layers = nn.Sequential(nn.Linear(in_features, hidden_features),
                                    nn.ReLU(),
                                    nn.Linear(hidden_features, hidden_features),
                                    nn.ReLU())
        self.log_s = nn.Linear(hidden_features, in_features)
        self.t = nn.Linear(hidden_features, in_features)

    def forward(self, x):
        hidden = self.layers(x)
        log_s = torch.tanh(self.log_s(hidden))
        t = self.t(hidden)
        return log_s, t
# flow
flow_list = []
for i in range(5):
        scale_translate_net = nn.Sequential(nn.Linear(x_dim, 256),
                                            nn.ReLU(),
                                            nn.Linear(256, 256),
                                            nn.ReLU(),
                                            nn.Linear(256, x_dim * 2))
        flow_list.append(AffineCoupling(in_features=2,
                                        scale_translate_net=ScaleTranslateNet(x_dim, 256),
                                        inverse_mask=(i%2!=0)))
        flow_list.append(BatchNorm1d(2))

f = FlowList(flow_list)
# inverse transformed distribution (z -> f^-1 -> x)
p = InverseTransformedDistribution(prior=prior, flow=f, var=["x"]).to(device)
print_latex(p)
model = ML(p, optimizer=optim.Adam, optimizer_params={"lr":1e-2})
print(model)
print_latex(model)
# plot training set
from sklearn import datasets
x = datasets.make_moons(n_samples=test_size, noise=0.1)[0].astype("float32")
noise = prior.sample(batch_n=test_size)["z"].data.cpu()
plot_samples(x, noise)
for epoch in range(epochs):
    x = datasets.make_moons(n_samples=batch_size, noise=0.1)[0].astype("float32")
    x = torch.tensor(x).to(device)
    loss = model.train({"x": x})

    if epoch % 500 == 0:
        print('Epoch: {} Test loss: {:.4f}'.format(epoch, loss))
        
        # samples
        samples = p.sample(batch_n=test_size)["x"].data.cpu()
        
        # inference
        _x = datasets.make_moons(n_samples=test_size, noise=0.1)[0].astype("float32")
        _x = torch.tensor(_x).to(device)
        noise = p.inference({"x": _x})["z"].data.cpu()
        
        plot_samples(samples, noise)
samples = p.sample(batch_n=test_size)["x"].data.cpu()

# inference
_x = datasets.make_moons(n_samples=test_size, noise=0.1)[0].astype("float32")
_x = torch.tensor(_x).to(device)
noise = p.inference({"x": _x})["z"].data.cpu()

plot_samples(samples, noise)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14343742066514345, 'fraction_numerical': 0.024625539477024627, 'mean_word_length': 2.6313364055299537, 'pattern_counts': {'":': 4, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 13, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '420378', 'license': 'MIT', 'n_tokens_mistral': 1429, 'n_tokens_neox': 1355, 'n_words': 255, 'path': 'examples/real_nvp_toy.ipynb', 'repo_name': 'MokkeMeguru/pixyz-test', 'size': '3939'}