<jupyter_start><jupyter_text>CIFAR-10 training First, change the optimizer structure using the MetaConv2D.<jupyter_code>class CIFAR10Loss:
    def __init__(self, training=True):
        dataset = datasets.CIFAR10(
            '/datasets/cifar10', train=True, download=False,
            transform=torchvision.transforms.ToTensor()
        )
        indices = list(range(len(dataset)))
        np.random.RandomState(10).shuffle(indices)
        if training:
            indices = indices[:len(indices) // 2]
        else:
            indices = indices[len(indices) // 2:]

        self.loader = torch.utils.data.DataLoader(
            dataset, batch_size=128,
            sampler=torch.utils.data.sampler.SubsetRandomSampler(indices))

        self.batches = []
        self.cur_batch = 0
        
    def sample(self):
        if self.cur_batch >= len(self.batches):
            self.batches = []
            self.cur_batch = 0
            for b in self.loader:
                self.batches.append(b)
        batch = self.batches[self.cur_batch]
        self.cur_batch += 1
        return batch

class CIFAR10Net(MetaModule):

    def __init__(self, layer_size=32, n_layers=3, **kwargs):
        super().__init__()


        # add linear layers 
        
        # for i in range(n_layers):
        #     self.layers[f'mat_{i}'] = MetaLinear(inp_size, layer_size)
        #     inp_size = layer_size

        self.layers = {}
        # Main layers (Convolutions + MaxPooling)
        in_channels = 3
        self.hidden_channels = [3,5,8,16]

        for i in range(1,n_layers+1):
            self.layers[f'conv_{i}'] = MetaConv2d(in_channels, self.hidden_channels[i], kernel_size=3)
            self.layers[f'norm_{i}'] = nn.BatchNorm2d(self.hidden_channels[i-1], affine=False, track_running_stats=False)
            in_channels = self.hidden_channels[i]

        self.activation = nn.ReLU()
        self.pool = nn.MaxPool2d(kernel_size=2,)
        

        # Linear Layers 
        self.layers['linear'] = MetaLinear(16*2*2 , layer_size)

        self.layers = nn.ModuleDict(self.layers)

        self.loss = nn.NLLLoss()

    def all_named_parameters(self):
        return [(k, v) for k, v in self.named_parameters()]
    
    def forward(self, loss):
        inp, out = loss.sample()
        inp = w(Variable(inp.view(inp.size()[0], 3,32,32)))    # reshapes from inp.shape
        out = w(Variable(out))


        for cur_layer in range(1,4):    #since the layers are defined from 1...
            # print("Shape before conv2D - kernel = 3 \t" + str(inp.shape))
            inp = self.layers[f'norm_{cur_layer}'](inp)    #affine = False makes the batch nor parameters not-learnable (keep it simple)
            inp = self.layers[f'conv_{cur_layer}'](inp)
            inp = self.pool(inp)
            inp = self.activation(inp)
            
            cur_layer += 1
            
        # print("Shape after CNN \t" + str(inp.shape))
        # inp = inp.view(-1, 16*2*2)  # flatten
        inp = torch.flatten(inp,start_dim=1)
        
        inp = F.log_softmax(self.layers['linear'](inp), dim=1)

        #debugging....
        # print(inp.shape,out.shape)
        l = self.loss(inp, out)

        return l<jupyter_output><empty_output><jupyter_text>After the optmizee is done, we re-build the optmizer with athe new architecture: 2 LSTMs in series - one forthe conv2Dand one for the lienear (i.e.) fully connected.<jupyter_code>class Optimizer(nn.Module):
    def __init__(self, preproc=False, hidden_sz=20, preproc_factor=10.0):
        super().__init__()
        self.hidden_sz = hidden_sz
        if preproc:
            self.recurs = nn.LSTMCell(2, hidden_sz)
        else:
            self.recurs = nn.LSTMCell(1, hidden_sz)
        self.recurs2 = nn.LSTMCell(hidden_sz, hidden_sz)
        self.output = nn.Linear(hidden_sz, 1)
        self.preproc = preproc
        self.preproc_factor = preproc_factor
        self.preproc_threshold = np.exp(-preproc_factor)

    def forward(self, inp, hidden, cell):
        if self.preproc:
            # Implement preproc described in Appendix A

            # Note: we do all this work on tensors, which means
            # the gradients won't propagate through inp. This
            # should be ok because the algorithm involves
            # making sure that inp is already detached.
            inp = inp.data
            inp2 = w(torch.zeros(inp.size()[0], 2))
            keep_grads = (torch.abs(inp) >= self.preproc_threshold).squeeze()
            inp2[:, 0][keep_grads] = (torch.log(torch.abs(inp[keep_grads]) + 1e-8) / self.preproc_factor).squeeze()
            inp2[:, 1][keep_grads] = torch.sign(inp[keep_grads]).squeeze()

            inp2[:, 0][~keep_grads] = -1
            inp2[:, 1][~keep_grads] = (float(np.exp(self.preproc_factor)) * inp[~keep_grads]).squeeze()
            inp = w(Variable(inp2))
        hidden0, cell0 = self.recurs(inp, (hidden[0], cell[0]))
        hidden1, cell1 = self.recurs2(hidden0, (hidden[1], cell[1]))
        
        return self.output(hidden1), (hidden0, hidden1), (cell0, cell1)<jupyter_output><empty_output><jupyter_text>Now let's try to fit CIFAR classifier.<jupyter_code>def get_fit_dict_test(n_tests, opt_dict, *args, **kwargs):
    opt = w(Optimizer(preproc=True))
    # print(opt_dict)
    # opt.load_state_dict(opt_dict)
    np.random.seed(0)
    return [do_fit(opt, *args, **kwargs) for _ in tqdm(range(N_TESTS), 'optimizer')]



def fit_normal(target_cls, target_to_opt, opt_class, n_tests=100, n_epochs=100, **kwargs):
    results = []
    for i in tqdm(range(n_tests), 'tests'):
        target = target_cls(training=False)
        optimizee = w(target_to_opt())
        optimizer = opt_class(optimizee.parameters(), **kwargs)
        total_loss = []
        for _ in range(n_epochs):
            loss = optimizee(target)
            
            total_loss.append(loss.data.cpu().numpy())
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        results.append(total_loss)

        
    return results<jupyter_output><empty_output><jupyter_text>Plotting<jupyter_code>import torch.optim as optim
import torchvision as tv

dataset = tv.datasets.CIFAR10(
            '/datasets/cifar10', train=True, download=True,
            transform=tv.transforms.ToTensor()
        )


device = torch.device('cpu')
mnist_optimizer = torch.load('mnist_model.pt', map_location=device)

NORMAL_OPTS = [(optim.Adam, {}), (optim.RMSprop, {}), (optim.SGD, {'momentum': 0.9}), (optim.SGD, {'nesterov': True, 'momentum': 0.9})]
OPT_NAMES = ['ADAM', 'RMSprop', 'SGD', 'NAG']
LEARNING_RATES = [0.01, 0.003, 0.03, 0.01]
N_TESTS = 50
N_EPOCHS = 50

fit_data = np.zeros((N_TESTS, N_EPOCHS, len(OPT_NAMES) + 1))
for i, ((opt, extra_kwargs), lr) in enumerate(zip(NORMAL_OPTS, LEARNING_RATES)):
    np.random.seed(0)
    fit_data[:, :, i] = np.array(fit_normal(CIFAR10Loss, CIFAR10Net, opt, lr=lr, n_tests=N_TESTS, n_epochs=N_EPOCHS, **extra_kwargs))


fit_data[:, :, len(OPT_NAMES)] = np.array(get_fit_dict_test(N_TESTS, mnist_optimizer, None, CIFAR10Loss, CIFAR10Net, N_TESTS, N_EPOCHS, N_EPOCHS, out_mul=0.1, should_train=False))
import pandas as pd
condition=OPT_NAMES + ['LSTM']
def to_long_form(data_arr, condition):
    records = []
    for i in range(data_arr.shape[0]):
        for j in range(data_arr.shape[1]):
            for k in range(data_arr.shape[2]):
                records.append((i, j, condition[k], data_arr[i, j, k]))
    df = pd.DataFrame.from_records(records, columns=['test', 'Steps', 'Optimizer', 'Loss'] )
    return df

COLORS = ["red", "blue", "green", "purple", "orange"]

df = to_long_form(fit_data, condition)
axis = sns.lineplot(data=df, x='Steps', y='Loss', hue='Optimizer', ci=95, palette=COLORS)
plt.gcf().set_figwidth(11.7)

plt.yscale('log')
plt.title("CIFAR10")
plt.show()
np.save('mnist-on-cifar.npy', fit_data)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13717819283190308, 'fraction_numerical': 0.019687026754164564, 'mean_word_length': 2.731167608286252, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 19, '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': '22731', 'license': 'MIT', 'n_tokens_mistral': 2892, 'n_tokens_neox': 2767, 'n_words': 532, 'path': 'plotting/plotting-cifar.ipynb', 'repo_name': 'Pluriscient/learn-to-learn', 'size': '7924'}