<jupyter_start><jupyter_text>Model Definition<jupyter_code>class MLP(nn.Module):
    def __init__(self, input_size, output_size, hidden_sizes=None):
        super(MLP, self).__init__()
        self.input_size = input_size
        self.hidden_sizes = hidden_sizes
        if output_size is not None:
            self.output_size = output_size
            self.squeeze_output = False
        else:
            self.output_size = 1
            self.squeeze_output = True
        self.act = F.relu
        # Define layers
        if hidden_sizes is None:
            # Linear model
            self.hidden_layers = []
            self.output_layer = nn.Linear(self.input_size, self.output_size)
        else:
            # Neural network
            features = zip([self.input_size] + hidden_sizes[:-1], hidden_sizes)
            self.hidden_layers = nn.ModuleList([nn.Linear(in_features, out_features) for in_features, out_features in features])
            self.output_layer = nn.Linear(hidden_sizes[-1], self.output_size)

    def forward(self, x):
        x = x.view(-1, self.input_size)
        h = x
        for layer in self.hidden_layers:
            h = self.act(layer(h))
        out = self.output_layer(h)
        if self.squeeze_output:
            out = torch.squeeze(out).view([-1])
        return out
# Model arguments
model_kwargs = dict(input_size=2, 
                    output_size=None, 
                    hidden_sizes=[128])
model1 = MLP(**model_kwargs)
model1 = model1.to(device)
optimizer1 = torch.optim.Adam(model1.parameters())

model2 = MLP(**model_kwargs)
model2 = model2.to(device)
optimizer2 = torchsso.optim.VOGN(model2, dataset_size=len(train_loader.dataset))<jupyter_output><empty_output><jupyter_text>Train Loop<jupyter_code>figpaths = []
i = 0  # iteration
# Run training
for epoch in range(epochs):

    model1.train()
    model2.train()

    for data, target in train_loader:

        data, target = data.to(device), target.to(device)

        def closure1():
            optimizer1.zero_grad()
            output = model1(data)
            loss = F.binary_cross_entropy_with_logits(output, target)
            loss.backward()
            return loss

        def closure2():
            optimizer2.zero_grad()
            output = model2(data)
            loss = F.binary_cross_entropy_with_logits(output, target)
            loss.backward()
            return loss, output

        loss1 = optimizer1.step(closure1)
        loss2, _ = optimizer2.step(closure2)

        if (i + 1) % plot_interval == 0:
            # Setup figures
            fig = plt.figure(figsize=(21, 6))
            gs = fig.add_gridspec(1, 3)

            # Decision boundary
            ax1 = fig.add_subplot(gs[0, 0])
            ax1.set_xlabel('Input 1')
            ax1.set_ylabel('Input 2')
            ax1.set_title(f'Iteration {i+1}')

            # Entropy (Adam)
            ax2 = fig.add_subplot(gs[0, 1])
            ax2.set_xlabel('Input 1')
            ax2.set_ylabel('Input 2')
            ax2.set_title(f'Entropy (Adam)')

            # Entropy (VOGN)
            ax3 = fig.add_subplot(gs[0, 2])
            ax3.set_xlabel('Input 1')
            ax3.set_ylabel('Input 2')
            ax3.set_title(f'Entropy (VOGN)')

            model1.eval()
            model2.eval()

            # (Adam)
            prob = torch.sigmoid(model1(data_meshgrid)).view(xx.shape)
            entropy = get_entropy(prob)
            pred = torch.round(prob).detach().cpu().numpy()

            plot = ax1.contour(xx, yy, pred, colors=['blue'], linewidths=[2])
            plot.collections[len(plot.collections)//2].set_label('Adam')
            im = ax2.pcolormesh(xx, yy, entropy)
            fig.colorbar(im, ax=ax2)

            # (VOGN) get MC samples
            prob, probs = optimizer2.prediction(data_meshgrid, keep_probs=True)
            prob = prob.view(xx.shape)
            entropy = get_entropy(prob)

            probs = probs[:n_samples_for_mcplot]
            preds = [torch.round(p).detach().cpu().numpy().reshape(xx.shape) for p in probs]
            for pred in preds:
                ax1.contour(xx, yy, pred, colors=['red'], alpha=0.01)
            im = ax3.pcolormesh(xx, yy, entropy)
            fig.colorbar(im, ax=ax3)

            # (VOGN) get mean prediction
            prob = optimizer2.prediction(data_meshgrid, mc=0).view(xx.shape)
            pred = torch.round(prob).detach().cpu().numpy()

            plot = ax1.contour(xx, yy, pred, colors=['red'], linewidths=[2])
            plot.collections[len(plot.collections)//2].set_label('VOGN')

            # plot samples
            for label, marker, color in zip([0, 1], ['o', 's'], ['white', 'gray']):
                _X = X[y == label]
                ax1.scatter(_X[:, 0], _X[:, 1], s=80, c=color, edgecolors='black', marker=marker)
                ax2.scatter(_X[:, 0], _X[:, 1], s=80, c=color, edgecolors='black', marker=marker)
                ax3.scatter(_X[:, 0], _X[:, 1], s=80, c=color, edgecolors='black', marker=marker)

            # save tmp figure
            ax1.grid(linestyle='--')
            ax2.grid(linestyle='--')
            ax3.grid(linestyle='--')
            ax1.set_yticks([-5, 0, 5, 10])
            ax2.set_yticks([-5, 0, 5, 10])
            ax3.set_yticks([-5, 0, 5, 10])
            ax1.legend(loc='lower right')
            ax1.set_aspect(0.8)
            plt.tight_layout()
            figname = f'iteration{i+1}.png'
            figpath = os.path.join(fig_dir, figname)
            if not os.path.isdir(fig_dir):
                os.makedirs(fig_dir)
            fig.savefig(figpath)
            plt.close(fig)
            figpaths.append(figpath)

        i += 1

    print(f'Train Epoch: {epoch+1}\tLoss(Adam): {loss1:.6f} Loss(VOGN): {loss2:.6f}')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1326442721791559, 'fraction_numerical': 0.023255813953488372, 'mean_word_length': 2.005175983436853, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '406000', 'license': 'MIT', 'n_tokens_mistral': 1989, 'n_tokens_neox': 1918, 'n_words': 350, 'path': 'notebooks/7.uncertainty_vogn_initial.ipynb', 'repo_name': 'rsilveira79/fermenting_gradients_2', 'size': '5805'}