<jupyter_start><jupyter_text>SGLD from here onwards<jupyter_code>@torch.enable_grad()
def gradient(x, y, params):
    params_ = params.clone().requires_grad_(True)
    loss = log_posterior(x, y, params_)
    grad, = torch.autograd.grad(loss, params_)
    return loss.detach().cpu().numpy(), grad
def step_size(n):
    return 1e-4 / (1 + n)**0.55
def sgld(n_epochs, data_batch_size):
    dataloader_train = DataLoader(CIFAR10_train, shuffle=True, batch_size=data_batch_size, num_workers=2)
    params = torch.cat([param.flatten() for param in model.parameters()]).detach()
    losses = []
    step = 0
    for _ in range(n_epochs):
        epoch_losses = []
        for x, y in tqdm(iter(dataloader_train)):
            x = x.to(device)
            y = y.to(device)

            eps = step_size(step)
            loss, grad = gradient(x, y, params)
            params = params + 0.5 * eps * grad #+ np.sqrt(eps) * torch.randn_like(params)
            step += 1
            epoch_losses.append(loss)
        
        losses.append(epoch_losses)
    
    param_samples = []
    
    iterator = iter(dataloader_train)
    for _ in range(100):
        x = x.to(device)
        y = y.to(device)
        
        eps = step_size(step)
        loss, grad = gradient(x, y, params)
        params = params + 0.5 * eps * grad + np.sqrt(eps) * torch.randn_like(params)
        param_samples.append(params)
        step += 1
        
    param_samples = torch.stack(param_samples)
    losses = np.array(losses)
    
    return param_samples, losses
accuracies, eces, logps = [], [], []

for i in range(5):
    param_samples, losses = sgld(n_epochs, data_batch_size)
    
    accuracy, ece, logp = evaluate(param_samples)
    
    accuracies.append(accuracy)
    eces.append(ece)
    logps.append(logp)
    
accuracies = np.array(accuracies)
eces = np.array(eces)
logps = np.array(logps)
SGLD_df = pd.DataFrame({"Accuracy": accuracies, "ECE": eces, "log predictive": logps})
SGLD_df
SGLD_df.describe()
def sgd(n_epochs, data_batch_size):
    
    dataloader_train = DataLoader(CIFAR10_train, shuffle=True, batch_size=data_batch_size, num_workers=2)
    model = LeNet5Cifar10(10).to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=1e-2)
    losses = []

    for i in range(n_epochs):
        for x, y in tqdm(iter(dataloader_train)):

            x = x.to(device)
            y = y.to(device)

            optimizer.zero_grad()

            out = model(x)

            l = F.cross_entropy(out, y, reduction="mean")

            l.backward()

            losses.append(l.detach().cpu().numpy())

            optimizer.step()

    losses = np.array(losses)
    
    return model, losses
accuracies, eces, logps = [], [], []

for i in range(5):
    model, losses = sgd(n_epochs, data_batch_size)
    params = model.parameters()
    params = torch.cat([param.flatten() for param in params]).detach()
    params = params.view(1, -1)
    accuracy, ece, logp = evaluate(params)
    
    accuracies.append(accuracy)
    eces.append(ece)
    logps.append(logp)
    
accuracies = np.array(accuracies)
eces = np.array(eces)
logps = np.array(logps)
SGD_df = pd.DataFrame({"Accuracy": accuracies, "ECE": eces, "log predictive": logps})
SGD_df
SGD_df.describe()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13274336283185842, 'fraction_numerical': 0.01007018614586512, 'mean_word_length': 2.7591743119266057, 'pattern_counts': {'":': 6, '<': 5, '<?xml version=': 0, '>': 5, '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': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '126856', 'license': 'MIT', 'n_tokens_mistral': 1205, 'n_tokens_neox': 1183, 'n_words': 229, 'path': 'notebooks/Bayesian Neural Networks CIFAR10.ipynb', 'repo_name': 'franciscovargas/ControlledFollmerDrift', 'size': '3277'}