<jupyter_start><jupyter_text>Bayesian Neural Network- Use an SGLD kernel and compute the accuracy on the test dataset while training- MNIST dataset<jupyter_code>from sgmcmcjax.models.bayesian_NN.NN_data import X_train, X_test, y_train, y_test
from sgmcmcjax.models.bayesian_NN.NN_model import init_network, loglikelihood, logprior, accuracy

from sgmcmcjax.kernels import build_sgld_kernel
from tqdm.auto import tqdm

batch_size = int(0.01*X_train.shape[0])
data = (X_train, y_train)

init_fn, my_kernel, get_params = build_sgld_kernel(5e-5, loglikelihood, logprior, data, batch_size)

my_kernel = jit(my_kernel)
# define the inital state
key = random.PRNGKey(10)
key, subkey1, subkey2 = random.split(key,3)
sizes = [784, 100, 10]
params_IC = init_network(subkey1, sizes)
state = init_fn(subkey2, params_IC)
%%time

Nsamples = 2000
samples = []
accuracy_list = []

for i in tqdm(range(Nsamples)):
    key, subkey = random.split(key)
    state = my_kernel(i, subkey, state)
    samples.append(get_params(state))
    if i%10==0:
        accuracy_list.append(accuracy(get_params(state), X_test, y_test))
    

plt.plot(accuracy_list[:])

print(f"Final accuracy: {100*accuracy_list[-8]:.1f}%")<jupyter_output>Final accuracy: 92.1%
{'fraction_non_alphanumeric': 0.1133768352365416, 'fraction_numerical': 0.02936378466557912, 'mean_word_length': 5.854748603351955, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, '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': '474383', 'license': 'Apache-2.0', 'n_tokens_mistral': 494, 'n_tokens_neox': 466, 'n_words': 103, 'path': 'docs/nbs/BNN.ipynb', 'repo_name': 'dfm/SGMCMCJax', 'size': '1226'}