<jupyter_start><jupyter_text>Introducing data loaders and batches<jupyter_code>import patches
import torch
import torch.utils as utils
import lettertask
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
import numpy as np
import plotnine as gg
import lazytools_sflippl as lazytools
import time
import pandas as pd
import mizani.transforms as transforms
cbm = lettertask.data.CompositionalBinaryModel(
    width=[5, 5],
    change_probability=[0.05, 0.5],
    samples=10000,
    seed=None
)
cts = patches.data.Contrastive1DTimeSeries(cbm.to_array())
cdl = utils.data.DataLoader(cts, batch_size=16, drop_last=True)
ce = patches.networks.LinearScaffold(1, data=cts)
criterion = patches.losses.ContrastiveLoss(loss=nn.MSELoss())
similarity = patches.losses._NegExp(loss=nn.MSELoss())
criterion(x)<jupyter_output><empty_output><jupyter_text>Seems to work!<jupyter_code>ideal = np.array([1]*5+[0]*5)/np.sqrt(5)
optimizer = optim.Adam(ce.parameters())
angles = []
loss_traj = []
running_loss = 0
for epoch in tqdm(range(5)):
    running_loss = 0
    for i, data in enumerate(cdl):
        if i<len(cdl):
            if i % 10 == 0:
                est = list(ce.parameters())[0].detach().numpy()
                angles.append(np.matmul(ideal, est.T)/np.sqrt(np.matmul(est, est.T)))
            optimizer.zero_grad()
            code = ce(data)
            loss = criterion(code)
            loss.backward()
            optimizer.step()
            running_loss += loss
            if i % 50 == 49:
                loss_traj.append(running_loss.detach().numpy()/50)
                running_loss = 0
(gg.ggplot(lazytools.array_to_dataframe(np.array(loss_traj)),
           gg.aes(x='dim0', y='array')) +
    gg.geom_line())
(gg.ggplot(lazytools.array_to_dataframe(np.array(angles)),
           gg.aes(x='dim0', y='array')) + 
    gg.geom_line())<jupyter_output><empty_output><jupyter_text>Exploring different batch sizes<jupyter_code>2**12
time.time()-start
loss_dfs = []
angle_dfs = []
batch_sizes = range(12)
for batch_size in batch_sizes:
    cdl = utils.data.DataLoader(cts, batch_size=2**batch_size, drop_last=True)
    ce = patches.networks.LinearScaffold(1, data=cts)
    optimizer = optim.Adam(ce.parameters(), lr=0.01)
    angles = []
    loss_traj = []
    running_loss = 0
    print('Starting batch size {}'.format(2**batch_size))
    start = time.time()
    while (time.time()-start) < 2:
        running_loss = 0
        for i, data in enumerate(cdl):
            if i<len(cdl):
                if i % np.ceil(10/(2**batch_size)) == 0:
                    est = list(ce.parameters())[0].detach().numpy()
                    angles.append(np.matmul(ideal, est.T)/np.sqrt(np.matmul(est, est.T)))
                optimizer.zero_grad()
                code = ce(data)
                loss = criterion(code)
                loss.backward()
                optimizer.step()
                running_loss += loss
                if i % np.ceil(50/(2**batch_size)) == np.ceil(50/(2**batch_size))-1:
                    loss_traj.append(running_loss.detach().numpy()/
                                     np.ceil(50/(2**batch_size)))
                    running_loss = 0
    stop = time.time()-start
    loss_df = lazytools.array_to_dataframe(np.array(loss_traj))\
                       .rename(columns={'dim0': 'steps', 'array': 'loss'})
    loss_df['steps'] *= stop/loss_df['steps'].max()
    loss_df['batch_size'] = np.array(2**batch_size)
    angle_df = lazytools.array_to_dataframe(np.array(angles))\
                        .rename(columns={'dim0': 'steps', 'array': 'angle'})
    angle_df['steps'] *= stop/angle_df['steps'].max()
    angle_df['batch_size'] = np.array(2**batch_size)
    loss_dfs.append(loss_df)
    angle_dfs.append(angle_df)
df_loss = pd.concat(loss_dfs)
df_angle = pd.concat(angle_dfs)
df_loss.groupby('batch_size')['steps'].max()
(gg.ggplot(df_loss[df_loss['steps']<=5],
           gg.aes(x='steps', y='loss', color='batch_size', group='batch_size')) + 
    gg.geom_line() + 
    gg.scale_color_continuous(trans=transforms.log2_trans) +
    gg.facet_wrap('batch_size'))
(gg.ggplot(df_angle[df_angle['steps']<=5],
           gg.aes(x='steps', y='angle', color='batch_size', group='batch_size')) + 
    gg.geom_line() + 
    gg.scale_color_continuous(trans=transforms.log2_trans) +
    gg.facet_wrap('batch_size'))<jupyter_output><empty_output><jupyter_text>Minibatches indeed appear optimal. Let's dive further into the area between 8 and 16. Clearly we should not overfit, but it is interesting to explore how smooth the angle can actually be. Also note that this does not disadvantage large batches since one step always represents about one second.<jupyter_code>del df_loss
del df_angle
loss_dfs = []
angle_dfs = []
batch_sizes = range(8, 16)
for batch_size in batch_sizes:
    cdl = utils.data.DataLoader(cts, batch_size=batch_size, drop_last=True)
    ce = patches.networks.LinearScaffold(1, data=cts)
    optimizer = optim.Adam(ce.parameters(), lr=0.01)
    angles = []
    loss_traj = []
    running_loss = 0
    print('Starting batch size {}'.format(batch_size))
    start = time.time()
    while (time.time()-start) < 2:
        running_loss = 0
        for i, data in enumerate(cdl):
            if i<len(cdl):
                if i % np.ceil(10/(2**batch_size)) == 0:
                    est = list(ce.parameters())[0].detach().numpy()
                    angles.append(np.matmul(ideal, est.T)/np.sqrt(np.matmul(est, est.T)))
                optimizer.zero_grad()
                code = ce(data)
                loss = criterion(code)
                loss.backward()
                optimizer.step()
                running_loss += loss
                if i % np.ceil(50/(batch_size)) == np.ceil(50/(batch_size))-1:
                    loss_traj.append(running_loss.detach().numpy()/
                                     np.ceil(50/(batch_size)))
                    running_loss = 0
    stop = time.time()-start
    loss_df = lazytools.array_to_dataframe(np.array(loss_traj))\
                       .rename(columns={'dim0': 'steps', 'array': 'loss'})
    loss_df['steps'] *= stop/loss_df['steps'].max()
    loss_df['batch_size'] = np.array(batch_size)
    angle_df = lazytools.array_to_dataframe(np.array(angles))\
                        .rename(columns={'dim0': 'steps', 'array': 'angle'})
    angle_df['steps'] *= stop/angle_df['steps'].max()
    angle_df['batch_size'] = np.array(batch_size)
    loss_dfs.append(loss_df)
    angle_dfs.append(angle_df)
df_loss = pd.concat(loss_dfs)
df_angle = pd.concat(angle_dfs)
(gg.ggplot(df_loss[df_loss['steps']<=5],
           gg.aes(x='steps', y='loss', color='batch_size', group='batch_size')) + 
    gg.geom_line() + 
    gg.scale_color_continuous() + 
    gg.facet_wrap('batch_size'))
(gg.ggplot(df_angle[df_angle['steps']<=5],
           gg.aes(x='steps', y='angle', color='batch_size', group='batch_size')) + 
    gg.geom_line() + 
    gg.scale_color_continuous() +
    gg.facet_wrap('batch_size'))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1496018202502844, 'fraction_numerical': 0.014931740614334471, 'mean_word_length': 2.8792057363485934, 'pattern_counts': {'":': 0, '<': 26, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '202229', 'license': 'MIT', 'n_tokens_mistral': 2536, 'n_tokens_neox': 2391, 'n_words': 402, 'path': 'notebooks/07-batches.ipynb', 'repo_name': 'sflippl/patches', 'size': '7032'}