<jupyter_start><jupyter_text>Making dataset<jupyter_code>ind = 1
T = True
S = []
B = []
while ind<20:

    ckpt = 'curves_mnist/LinearOneLayer/LongTraining/curve'+str(ind)+'/checkpoint-30.pt'
    checkpoint = torch.load(ckpt)
    m.load_state_dict(checkpoint['model_state'])
    
    S.append(samples(m))
    B.append(list(m.parameters())[-1].data.numpy())
    ind+=1
    
S = np.concatenate(S)
# len(S)/512
from torch.optim.lr_scheduler import StepLR

def train(flow, model1, model2, epochs=20, lr=1e-5, batchsize=512, step=10, test_every=20, start_epoch=1):
    flow.cuda();
    flow.train();
    trainable_parametrs = filter(lambda param: param.requires_grad, flow.parameters()) # list of all trainable parameters in a flow
    optimizer = torch.optim.Adam(trainable_parametrs, lr=lr, weight_decay=1e-4)# choose an optimizer, use module torch.optim
    scheduler = StepLR(optimizer, step_size=20, gamma=0.9)# decrease learning rate by torch.optim.lr_scheduler

    for epoch in range(start_epoch, epochs+1):  
        scheduler.step()
        flow.train()
        
        t = time.time()
        total_loss = 0

        for X in iterate_minibatches(S, batchsize):

            loss = -flow.log_prob(X).mean() #  compute the maximum-likelihood loss

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss+=loss
        
        total_loss /=(len(S)//batchsize)
        if epoch % step == 0:
            print('epoch %s:' % epoch, 'loss = %.3f' % total_loss, 'time = %.2f' %  (time.time()-t))
        if epoch % test_every == 0:
            test_flow(S, model1, model2, flow, N=2000)
        if epoch % 50 ==0:
            utils.save_checkpoint(
                'flows_checkpoints/RNVPjn',
                epoch,
                model_state=flow.state_dict(),
                optimizer_state=None
            )

# os.makedirs('flows_checkpoints/RNVPjn', exist_ok=True)<jupyter_output><empty_output><jupyter_text>Generator<jupyter_code>def iterate_minibatches(train_data, batchsize):
    indices = np.random.permutation(np.arange(len(train_data)))
    for start in range(0, len(indices), batchsize):
        ix = indices[start: start + batchsize]
        yield torch.FloatTensor(train_data[ix]).cuda()
from torch import distributions
prior = distributions.MultivariateNormal(torch.zeros(795), torch.eye(795))
N = 2
X = torch.FloatTensor(S[:N]).cuda()
# X.shape
def get_model(W, B):
    model_sampled = architecture.base(num_classes=10, **architecture.kwargs)
    model_samples = W.cpu().data.numpy()
    SIZE = 2000

    offset = 0
    for parameter in list(model_sampled.parameters())[:-1]:
        size = int(np.prod(parameter.size())/SIZE)
        value = model_samples[:, offset:offset+size]
        if size==10 or size==1:
             value = value.T
#         print(value.shape)
        value = value.reshape(parameter.size())
#         print(value.shape)
        parameter.data.copy_(torch.from_numpy(value))
        offset += size

    list(model_sampled.parameters())[-1].data.copy_(torch.from_numpy(B.mean(0)))
    
    return model_sampled

def test(model):
    criterion = F.cross_entropy
    regularizer = None 
    train_res = utils.test(loaders['train'], model, criterion, regularizer)
    test_res = utils.test(loaders['test'], model, criterion, regularizer) 
    print(train_res)
    print(test_res)

def test_flow(S, model1, model2, flow, N=2000):
    
    rcParams['figure.figsize'] = 12, 10
    rcParams['figure.dpi'] = 100

    flow.cuda()
    flow.eval()
    print('copmuting samples...')
    X = torch.FloatTensor(S[:N]).cuda()
    X_sample = X.data.cpu().numpy()
    X_prior = prior.sample((N,)).cpu().data.numpy()
    X_flow = flow.sample(N, ).data.cpu().numpy()
    X_sample_prior = flow.f(torch.FloatTensor(X_sample).cuda())[0].data.cpu().numpy()
    
    print('drawing...')
    i, j = 500, -1
    fig, axes = plt.subplots(2, 2,)
    axes[0, 0].set_title('Samples')
    axes[0, 0].scatter(X_sample[:, i], X_sample[:, j])
    axes[0, 1].set_title('Prior')
    axes[0, 1].scatter(X_prior[:, i], X_prior[:, j])
    axes[1, 0].set_title('Flow sampling')
    axes[1, 0].scatter(X_flow[:, i], X_flow[:, j])
    axes[1, 1].set_title('Map from samples to prior')
    axes[1, 1].scatter(X_sample_prior[:, i], X_sample_prior[:, j]) 
    plt.show()
    
    print('computing Arc model...')
    W1 = samples(model1)
    W2 = samples(model2)
    
#     flow.cpu()
    W_pre = 1/np.sqrt(2)*flow.f(torch.FloatTensor(W1).cuda())[0]+1/np.sqrt(2)*flow.f(torch.FloatTensor(W2).cuda())[0]
    W = flow.g(W_pre)
    B = []
    B.append(list(model1.parameters())[-1].data.cpu().numpy())
    B.append(list(model2.parameters())[-1].data.cpu().numpy())
    B = np.array(B)
    
    model_sampled = get_model(W, B)
    test(model_sampled)
    
    if N==2000:
        B = []
        B.append(flow.b2.data.cpu().numpy())
        B = np.array(B)
        print('computing Sampling from flow model...')
        X_flow = flow.sample(N, ).data.cpu()
        
        model_flow = get_model(X_flow, B)
        test(model_flow)
#         return model_sampled, model_flow
    
#     return model_sampled
    
class BatchNorm(nn.Module):
    def __init__(self, dim_in):
        super(BatchNorm, self).__init__()

        self.mu = torch.zeros(dim_in).cuda()
        self.sig2 = torch.zeros(dim_in).cuda()+0.1
        self.momentum=0.1
        
    def forward(self, x):

        if self.training:
            mu = x.mean(0)
            sig2 = (x-mu).pow(2).mean(0)  
            x = (x-mu)/(sig2 + 1.0e-6 ).sqrt()
            self.mu = self.momentum*mu + (1-self.momentum)*self.mu
            self.sig2 =  self.momentum*sig2 + (1-self.momentum)*self.sig2
            return x, sig2+1.0e-6 
        else: 
            x = (x-self.mu)/(self.sig2+1.0e-6).sqrt()
            return x, self.sig2+1.0e-6<jupyter_output><empty_output><jupyter_text>Flow Generator<jupyter_code># class SNet(nn.Module):
#     def __init__(self, dim_in, dim_middle):
#         super(SNet, self).__init__()
#         self.fc = nn.Sequential(
#             nn.Linear(dim_in, dim_middle),
#             nn.LeakyReLU(),
#             nn.Linear(dim_middle, dim_middle),
#             nn.LeakyReLU(),
#             nn.Linear(dim_middle, dim_in)
#         )
        
#     def forward(self, x):   
#         x = self.fc(x)
#         return x
    
# class TNet(nn.Module):
#     def __init__(self, dim_in, dim_middle):
#         super(TNet, self).__init__()
#         self.fc = nn.Sequential(
#             nn.Linear(dim_in, dim_middle),
#             nn.LeakyReLU(),
#             nn.Linear(dim_middle, dim_middle),
#             nn.LeakyReLU(),
#             nn.Linear(dim_middle, dim_in),
#         )
        
#     def forward(self, x):   
#         x = self.fc(x)
#         return x
    
# # class SNetF(nn.Module):
# #     def __init__(self, dim_in, dim_middle):
# #         super(SNetF, self).__init__()
# #         self.fc = nn.Sequential(
# #             nn.Linear(dim_in, dim_middle),
# #             nn.LeakyReLU(),
# #             nn.Linear(dim_middle, dim_middle),
# #             nn.LeakyReLU(),
# #             nn.Linear(dim_middle, dim_in)
# #         )
        
# #     def forward(self, x):   
# #         x = self.fc(x)
# #         return x
    
# # class TNetF(nn.Module):
# #     def __init__(self, dim_in, dim_middle):
# #         super(TNetF, self).__init__()
# #         self.fc = nn.Sequential(
# #             nn.Linear(dim_in, dim_middle),
# #             nn.LeakyReLU(),
# #             nn.Linear(dim_middle, dim_middle),
# #             nn.LeakyReLU(),
# #             nn.Linear(dim_middle, dim_in),
# #         )
        
# #     def forward(self, x):   
# #         x = self.fc(x)
# #         return x
    
    
class SNet(nn.Module):
    def __init__(self, dim_in, dim_middle):
        super(SNet, self).__init__()
        affine = True
        self.h = nn.Tanh() #nn.LeakyReLU() #nn.Tanh()
        self.fc = nn.Sequential(
            nn.Linear(dim_in, dim_middle),
            self.h,
            nn.BatchNorm1d(dim_middle, affine=affine),
            nn.Linear(dim_middle, dim_middle),
            self.h,
            nn.BatchNorm1d(dim_middle, affine=affine),
            nn.Linear(dim_middle, dim_in)
        )
        
    def forward(self, x):   
        x = self.fc(x)
#         x = torch.clamp(x, -1, 1) 
        return x
    
class TNet(nn.Module):
    def __init__(self,  dim_in, dim_middle):
        super(TNet, self).__init__()
        affine = True
        self.h = nn.Tanh()
        self.fc = nn.Sequential(
            nn.Linear(dim_in, dim_middle),
            self.h,
            nn.BatchNorm1d(dim_middle, affine=affine),
            nn.Linear(dim_middle, dim_middle),
            self.h,
            nn.BatchNorm1d(dim_middle, affine=affine),
            nn.Linear(dim_middle, dim_in),
        )
        
    def forward(self, x):   
        x = self.fc(x)
        return x

import numpy as np
N_layers = 7
n_dim = 795
onezero = [0, 1]*n_dim
masks = torch.Tensor([[onezero[:n_dim], onezero[1:n_dim+1]]]*N_layers)
masks = masks.view(2*N_layers, -1)
len(masks)
class RealNVP(nn.Module):
    def __init__(self, mask, prior):
        super(RealNVP, self).__init__()
        
        # Create a flow
        # nets:  a function that return a pytocrn neurel network e.g., nn.Sequential, s = nets(), s: dim(X) -> dim(X)
        # nett:  a function that return a pytocrn neurel network e.g., nn.Sequential, t = nett(), t: dim(X) -> dim(X)
        # mask:  a torch.Tensor of size #number_of_coupling_layers x #dim(X)
        # prior: an object from torch.distributions e.g., torch.distributions.MultivariateNormal
        
        self.prior = prior 
        self.mask = nn.Parameter(mask, requires_grad=False)
        self.t = torch.nn.ModuleList([SNet(dim_in=795, dim_middle=795*2) for _ in range(len(mask))])
        self.s = torch.nn.ModuleList([TNet(dim_in=795, dim_middle=795*2) for _ in range(len(mask))])
        self.b = torch.nn.ModuleList([BatchNorm(dim_in=795) for _ in range(len(mask))])
        self.batch_norm = True
        self.verbose = False
        
        # for prediction
        self.relu = nn.ReLU()
        data_b2 = torch.FloatTensor(np.random.normal(size=(10, ))).cuda()
        self.b2 = nn.Parameter(data = data_b2, requires_grad=True)
        
    def g(self, z):
        # Compute and return g(z) = x, 
        #    where self.mask[i], self.t[i], self.s[i] define a i-th masked coupling layer   
        # z: a torch.Tensor of shape batchSize x 1 x dim(X)
        # return x: a torch.Tensor of shape batchSize x 1 x dim(X)
        for i, (s, t, b) in enumerate(zip(reversed(self.s), reversed(self.t), reversed(self.b))):
            m = self.mask[-i-1]
#             print('i', i, 'm', m)
            
            if self.verbose:
                print('z1', z)
            z = (m*z +(1-m)*(z-t(m*z))*(-s(m*z)).exp()).detach()
#             print('z1', z)
            if self.batch_norm:
                z = (z*(b.sig2+1.0e-6).sqrt()+b.mu).detach()
            if self.verbose:
                print('z2', z)
        
        x = z
        return x

    def f(self, x):
        # Compute f(x) = z and log_det_Jakobian of f, 
        #    where self.mask[i], self.t[i], self.s[i] define a i-th masked coupling layer   
        # x: a torch.Tensor, of shape batchSize x dim(X), is a datapoint
        # return z: a torch.Tensor of shape batchSize x dim(X), a hidden representations
        # return log_det_J: a torch.Tensor of len batchSize
        
        z = x
        log_det_J = 0
        for s, t, m, b in zip(self.s, self.t, self.mask, self.b):

            if self.batch_norm:
                z, sig2 = b(z)
#             print(m.shape, z.shape)
            s_res = s(m*z)
            z = m*z +(1-m)*(z*s_res.exp()+t(m*z))
           
            if self.batch_norm:
                log_det_J += ((1-m)*s_res-0.5*sig2.log()).sum(-1)
            else:
                log_det_J += ((1-m)*s_res).sum(-1)
        
        return z, log_det_J
    
    def log_prob(self,x):
        # Compute and return log p(x)
        # using the change of variable formula and log_det_J computed by f
        # return logp: torch.Tensor of len batchSize
        z, log_det_J = self.f(x)
        
        logp = -0.5*np.log(np.pi*2)-0.5*z.pow(2)
        logp = logp.sum(-1)
#         logp = self.prior.log_prob(z)
#         logp = self.prior.log_prob(z.cpu()).cuda()
#         shape = torch.Size((K, self.in_dim))
#         logp = torch.cuda.FloatTensor(x.shape[0])
#         self.prior.log_prob(z.cpu(), out=logp)
        
#         logp = self.prior.log_prob(z.cpu()).cuda()
#         print('logp', logp.shape)
        
        return logp+log_det_J
        
    def sample(self, K): 
        # Draw and return batchSize samples from flow using implementation of g
        # return x: torch.Tensor of shape batchSize x 1 x dim(X)
        shape = torch.Size((K, 795))
        e = torch.cuda.FloatTensor(shape)
        torch.randn(shape, out=e)   
        x = self.g(e)
        
#         z = self.prior.sample((K, )).cuda()
#         x = self.g(z)
        
        return x
    
    def forward(self, x, K=2000):
        
        x = x.view(x.size(0), -1)
        W = self.sample(K)
        W1 = W[:, :784]
        b1 = W[:, 784:785]
        W2 =  W[:, 785:].transpose(0,1)
        
        if self.verbose:
            print('x', x.shape)
            print('W1', W1.shape)
            print('b1', b1.shape)
            
        x = F.linear(x, W1, b1[:, 0])
        
        if self.verbose:
            print('x', x.shape)
            print(x.shape)
            
        x = self.relu(x)

        x = F.linear(x, W2, self.b2) 
#         x = 1/K*x
        
        if self.verbose:
            print(x.shape)
            
        return x
        
def train(flow, model1, model2, epochs=20, lr=1e-5, batchsize=512, step=10, test_every=20, start_epoch=1):
    flow.cuda();
    flow.train();
    trainable_parametrs = filter(lambda param: param.requires_grad, flow.parameters()) # list of all trainable parameters in a flow
    optimizer = torch.optim.Adam(trainable_parametrs, lr=lr, weight_decay=0)# choose an optimizer, use module torch.optim
    scheduler = StepLR(optimizer, step_size=20, gamma=0.9)# decrease learning rate by torch.optim.lr_scheduler
    regularizer = None
    criterion = F.cross_entropy
    
    trainable_parametrs = filter(lambda param: param.requires_grad, flow.parameters())
    optimizer2 = torch.optim.Adam(trainable_parametrs, lr=1e-3, weight_decay=0)# choos
    
    columns = ['ep', 'tr_loss', 'tr_acc', 'te_nll', 'te_acc', 'time']
    
    for epoch in range(start_epoch, epochs+1):  
        scheduler.step()
        flow.train()
        
        t = time.time()
        total_loss = 0

        for X in iterate_minibatches(S, batchsize):

            loss = -flow.log_prob(X).mean() #  compute the maximum-likelihood loss

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss+=loss
        
        total_loss /=(len(S)//batchsize)
        if epoch % step == 0:
            print('epoch %s:' % epoch, 'loss = %.3f' % total_loss, 'time = %.2f' %  (time.time()-t))
        if epoch % 50 ==0:
            utils.save_checkpoint(
                'flows_checkpoints/RNVPjn',
                epoch,
                model_state=flow.state_dict(),
                optimizer_state=None
            )
        if epoch % 10 == 0:
            time_ep = time.time()
            train_res = utils.train(loaders['train'], flow, optimizer2, criterion, regularizer, cuda=True)
            test_res = utils.test(loaders['test'], flow, criterion, regularizer, cuda=True)
            time_ep = time.time() - time_ep
            values = [epoch, train_res['loss'], train_res['accuracy'], test_res['nll'],
                      test_res['accuracy'], time_ep]

            table = tabulate.tabulate([values], columns, tablefmt='simple', floatfmt='9.4f')
            print(table)
        if epoch % test_every == 0:
            test_flow(S, model1, model2, flow, N=2000)

# def train(flow, model1, model2, epochs=20, lr=1e-5, batchsize=512, step=10, test_every=20, start_epoch=1):
#     flow.cuda();
#     flow.train();
#     trainable_parametrs = filter(lambda param: param.requires_grad, flow.parameters()) # list of all trainable parameters in a flow
#     optimizer = torch.optim.Adam(trainable_parametrs, lr=lr, weight_decay=0)# choose an optimizer, use module torch.optim
#     scheduler = StepLR(optimizer, step_size=20, gamma=0.9)# decrease learning rate by torch.optim.lr_scheduler
#     regularizer = None
#     criterion = F.cross_entropy
    
#     trainable_parametrs = filter(lambda param: param.requires_grad, flow.parameters())
#     optimizer2 = torch.optim.Adam(trainable_parametrs, lr=1e-3, weight_decay=0)# choos
    
#     columns = ['ep', 'tr_loss', 'tr_acc', 'te_nll', 'te_acc', 'time']
#     train_loader = loaders['train']
    
#     for epoch in range(start_epoch, epochs+1):  
#         scheduler.step()
#         flow.train()
        
#         t = time.time()
#         total_loss = 0

#         for X, (input, target) in zip(iterate_minibatches(S, batchsize), train_loader):
            
#             input = input.cuda(non_blocking=True)
#             target = target.cuda(non_blocking=True)

#             loss = -flow.log_prob(X).mean() #  compute the maximum-likelihood loss
            
#             output = flow(input)
#             cross_entr = 10*criterion(output, target)
#             loss +=cross_entr
            
#             optimizer.zero_grad()
#             loss.backward()
#             optimizer.step()
#             total_loss+=loss
        
        
#         total_loss /=(len(S)//batchsize)
#         if epoch % step == 0:
#             print('epoch %s:' % epoch, 'loss = %.3f' % total_loss, 'loss = %.3f' % cross_entr, 'time = %.2f' %  (time.time()-t))
#         if epoch % test_every == 0:
#             test_flow(S, model1, model2, flow, N=2000)
#         if epoch % 50 ==0:
#             utils.save_checkpoint(
#                 'flows_checkpoints/RNVPjn',
#                 epoch,
#                 model_state=flow.state_dict(),
#                 optimizer_state=None
#             )
#         if epoch % 6 == 0:
#             test_res = utils.test(loaders['test'], flow, criterion, regularizer, cuda=True)
#             print(test_res)

flow = RealNVP(masks, prior)
flow.cuda();
i, j = 1, 500
flow.eval()
flow.g(flow.f(X)[0])[i][j], X[i][j]
# checkpoint = torch.load('flows_checkpoints/RNVPjn/checkpoint-50.pt')
# flow.load_state_dict(checkpoint['model_state'])
# train(flow, model1, model2, epochs=300, lr=1e-4, batchsize=1024, step=1, test_every=10)
train(flow, model1, model2, epochs=1000, lr=1e-4, batchsize=1024, step=1, test_every=10)
test_flow(S, model1, model2, flow, N=2000)<jupyter_output>copmuting samples...
drawing...
{'fraction_non_alphanumeric': 0.14571458729153472, 'fraction_numerical': 0.020371543170783196, 'mean_word_length': 2.118150403159454, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 14, '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': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '555062', 'license': 'BSD-2-Clause', 'n_tokens_mistral': 6888, 'n_tokens_neox': 6605, 'n_words': 1299, 'path': 'experiments/RNVP_LinearOneLayer.ipynb', 'repo_name': 'avecplezir/dnn-mode-connectivity', 'size': '18948'}