added data

data/mlp.py

@@ -0,0 +1,242 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.lines import Line2D
+import math
+
+def get_lr(optimizer):
+    for param_group in optimizer.param_groups:
+        return param_group['lr']
+
+class MLP(nn.Module):
+    def __init__(self, f_i: int, f_o: int, act_fn: object = nn.SELU, f=[], insert_in=[4], freq_encoding=False):
+        
+        super().__init__()
+
+        self.insert_in = insert_in
+        self.layers = nn.ModuleList()
+        self.act = act_fn()
+        self.freq_encoding = freq_encoding
+        f_in = f_i
+        for f_cntr, f_oo in enumerate(f):
+            if f_cntr in insert_in:
+                f_oo -= f_i
+            self.layers.append(nn.Linear(f_in, f_oo))
+            if f_cntr in insert_in:
+                f_in = f_oo+f_i
+            else:
+                f_in = f_oo
+
+        self.layers.append(nn.Linear(f_in, f_o))
+
+    def forward(self, x):
+
+        res = []
+        if self.freq_encoding:
+            
+            yc = x[:,-1:]
+            pi = torch.acos(torch.zeros(1)).item()
+            x = torch.cat((torch.cos(2.*pi*yc),
+                           torch.sin(2.*pi*yc),
+                           x[:,:-1]), axis=-1)
+        
+        inp = x
+        
+        for m_ind, m in enumerate(self.layers[:-1]):
+            x = m(x)
+           
+            if m_ind in self.insert_in:
+                x = torch.cat((inp,x), axis=1)
+
+            for r in res:
+                x += r
+                
+            x = self.act(x)
+            
+            res.append(x)
+
+        x = self.layers[-1](x)
+            
+        return x
+
+def one_epoch_mlp(mlp, epoch, loader, optimizer, device, is_train=False):
+    running_loss = 0.
+    counter = 1
+    loss_fn = torch.nn.L1Loss() #reduction="none")
+    
+    if is_train:
+        torch.set_grad_enabled(True)
+    else:
+        torch.set_grad_enabled(False)
+        
+    for i, data in enumerate(loader):
+        if 1==1: #is_train or i%10==0:
+            optimizer.zero_grad()
+
+            x = data[0].to(device)
+            y = data[1].to(device)
+            T = mlp(x.to(device))
+
+            loss  = loss_fn(y.to(device), T)
+            
+            if is_train:
+                loss.backward()
+                optimizer.step()
+
+            running_loss += loss.item()
+            counter += 1
+
+            if 1==2: #i % 1000 == 0:
+                print(epoch, i, running_loss/counter)
+    return running_loss/counter
+
+def one_epoch_mlp_lbfgs(mlp, epoch, loader, optimizer, device, is_train=False):
+    running_loss = 0.
+    counter = 1
+    loss_fn = torch.nn.L1Loss() #reduction="none")
+    
+    if is_train:
+        torch.set_grad_enabled(True)
+    else:
+        torch.set_grad_enabled(False)
+        
+    for i, data in enumerate(loader):
+        x = data[0].to(device)
+        y = data[1].to(device)
+        def closure():
+            optimizer.zero_grad()
+            
+            T = mlp(x.to(device))
+            loss  = loss_fn(y.to(device), T)
+                
+            loss.backward()
+            return loss
+
+        if is_train:
+            optimizer.step(closure)
+
+        T     = mlp(data[0].to(device))
+        loss  = loss_fn(data[1].to(device), T)
+
+        running_loss += loss.item()
+        counter += 1
+
+    return running_loss/counter
+
+def exists(val):
+    return val is not None
+
+def cast_tuple(val, repeat = 1):
+    return val if isinstance(val, tuple) else ((val,) * repeat)
+
+# sin activation
+class Sine(nn.Module):
+    def __init__(self, w0 = 1.):
+        super().__init__()
+        self.w0 = w0
+    def forward(self, x):
+        return torch.sin(self.w0 * x)
+
+# siren layer
+class Siren(nn.Module):
+    def __init__(
+        self,
+        dim_in,
+        dim_out,
+        w0 = 1.,
+        c = 6.,
+        is_first = False,
+        use_bias = True,
+        activation = None,
+        dropout = 0.
+    ):
+        super().__init__()
+        self.dim_in = dim_in
+        self.is_first = is_first
+
+        weight = torch.zeros(dim_out, dim_in)
+        bias = torch.zeros(dim_out) if use_bias else None
+        self.init_(weight, bias, c = c, w0 = w0)
+
+        self.weight = nn.Parameter(weight)
+        self.bias = nn.Parameter(bias) if use_bias else None
+        self.activation = Sine(w0) if activation is None else activation
+        self.dropout = nn.Dropout(dropout)
+
+    def init_(self, weight, bias, c, w0):
+        dim = self.dim_in
+
+        w_std = (1 / dim) if self.is_first else (math.sqrt(c / dim) / w0)
+        weight.uniform_(-w_std, w_std)
+
+        if exists(bias):
+            bias.uniform_(-w_std, w_std)
+
+    def forward(self, x):
+        out = F.linear(x, self.weight, self.bias)
+        out = self.activation(out)
+        out = self.dropout(out)
+        return out
+
+class SirenMLP(nn.Module):
+
+    def __init__(self,
+                 n_in,
+                 n_out,
+                 n_hidden,
+                 device,
+                 num_layers,
+                 w0 = 30.,
+                 w0_initial = 30.,
+                 use_bias = True,
+                 final_activation = None,
+                 dropout = 0.,
+                 context_params = None
+                ):
+        super().__init__() #SirenCavia, self
+
+        self.device = device
+        self.num_layers = num_layers
+        self.n_hidden = n_hidden
+        self.dropout = nn.Dropout(dropout)
+        
+        # siren layers
+        self.siren_layers = nn.ModuleList()
+        
+        for ind in range(num_layers):
+            is_first = ind == 0
+            layer_w0 = w0_initial if is_first else w0
+            layer_dim_in = n_in if is_first else n_hidden
+
+            layer = Siren(
+                dim_in = layer_dim_in,
+                dim_out = n_hidden,
+                w0 = layer_w0,
+                use_bias = use_bias,
+                is_first = is_first,
+                dropout = dropout
+            )
+
+            self.siren_layers.append(layer)
+
+        final_activation = nn.Identity() if not exists(final_activation) else final_activation
+        self.siren_layers.append(Siren(dim_in = n_hidden, dim_out = n_out, w0 = w0, 
+                                      use_bias = use_bias, activation = final_activation))
+
+        
+    def forward(self, x):
+        
+        res = []
+        for k in range(len(self.siren_layers) - 1):
+            x = self.siren_layers[k](x) 
+            for r in res:
+                x = 0.5*(x+r)
+            res.append(x)
+            x = self.dropout(x)
+            
+        x = self.siren_layers[-1](x)
+
+        return x

data/train_profiles_mlp.py

@@ -0,0 +1,152 @@
+import glob, os, sys
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+import torch
+from mlp import *
+import argparse
+from datasetio import *
+from torch.utils.data import TensorDataset
+
+import copy
+import pickle
+import time
+
+# In[ ]:
+
+
+data_dir = "/plp_scr1/agar_sh/data/TPH/"
+nn_dir = "/plp_user/agar_sh/PBML/pytorch/TPH/MLP/trained_networks/"
+
+
+# In[ ]:
+
+
+run_cell = True
+if run_cell:
+    parser = argparse.ArgumentParser(description='Train mlp')
+    parser.add_argument("-gpu", "--gpu_number", type=int, help="specify gpu number")
+    parser.add_argument("-a", "--act_fn", type=str, help ="activation function")
+    parser.add_argument("-l", "--num_layers", type=int, help ="activation function")
+    parser.add_argument("-f", "--f_h", type=int, help ="filters")
+    
+    args = parser.parse_args()
+    
+    gpu_number = args.gpu_number
+    act_fn = args.act_fn
+    num_layers = args.num_layers
+    f_h = [args.f_h]
+
+else:
+    f_h = [128]
+    num_layers = 3
+    act_fn = "selu"
+    gpu_number = 6
+
+
+# In[ ]:
+
+
+f_nn = "mlp_profile_pointwise_" + str(f_h) + "_" + str(num_layers) + "_" + act_fn
+
+if not os.path.isdir(nn_dir + f_nn):
+    os.mkdir(nn_dir + f_nn)
+
+device = torch.device("cuda:" + str(gpu_number)) if torch.cuda.is_available() else torch.device("cpu")
+
+epoch         = 0
+start_lr      = 1e-3
+milestones    = [10, 20, 30, 40, 50, 60]
+epochs        = 70
+best_vloss    = 1e+16
+batch_size    = 32
+
+f_i = 4
+f_o = 1
+
+mlp = MLP(f_i=f_i, f_o=f_o, f=f_h*num_layers, insert_in=[num_layers-1]).double().to(device)
+torch.compile(mlp)
+print(mlp)
+
+nn_dir = nn_dir + f_nn + "/"
+with open(nn_dir + "mlp.txt", 'w') as writer:
+    writer.write('Epoch, train loss, val loss, learning rate \n')
+
+
+# In[ ]:
+
+
+dataset = {}
+loader = {}
+batches = {}
+pre = "/plp_user/agar_sh/PBML/pytorch/TPH/MLP/profiles/"
+with open(pre + 'x_pointwise.pkl', 'rb') as file: 
+    x_pointwise = pickle.load(file) 
+with open(pre + 'y_pointwise.pkl', 'rb') as file: 
+    y_pointwise = pickle.load(file) 
+
+for an in ["train", "cv"]:
+    
+    '''
+    inds = x_pointwise[an][:,-1]>0.8
+    x_ = x_pointwise[an][inds,:]
+    y_ = y_pointwise[an][inds,:]
+
+    repeats = int(np.ceil(x_pointwise[an].shape[0]/(batch_size*inds.shape[0])))
+    
+    x_pointwise[an] = np.concatenate((x_pointwise[an], 
+                                      np.repeat(x_, repeats, axis=0)
+                                     ), axis=0)
+    y_pointwise[an] = np.concatenate((y_pointwise[an], 
+                                      np.repeat(y_, repeats, axis=0)
+                                     ), axis=0)
+    '''
+    
+    dataset[an] = TensorDataset(torch.tensor(x_pointwise[an], dtype=torch.float64),
+                                torch.tensor(y_pointwise[an], dtype=torch.float64)) #.view(-1,2))
+    batches[an] = int(len(dataset[an])/batch_size)
+    loader[an]  = DataLoader(dataset[an], batch_size=batch_size, shuffle=True)
+    print(an, len(dataset[an]), batches[an])
+
+
+# In[ ]:
+
+
+optimizer = torch.optim.Adam([
+                                {"params": mlp.parameters(),
+                                 "lr": start_lr,
+                                 "weight_decay": 5e-4
+                                }
+                              ])
+ 
+scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=milestones, gamma=0.5)
+
+
+# In[ ]:
+
+
+for epoch in range(epochs):
+    t0 = time.time()
+    mlp.train(True)
+    avg_loss = one_epoch_mlp(mlp, epoch, loader["train"], optimizer, device, is_train=True)
+    
+    mlp.eval()
+    avg_vloss = one_epoch_mlp(mlp, epoch, loader["cv"], optimizer, device, is_train=False)
+    
+    print("-------------------------------------------")
+    print(epoch, "train: ", avg_loss, get_lr(optimizer))
+    print(epoch, "cv: ", avg_vloss)
+    print("took: " + str(time.time() - t0))
+    print("-------------------------------------------")
+    
+    if avg_vloss < best_vloss:
+        best_vloss = avg_vloss
+        torch.save(mlp.state_dict(), nn_dir + "mlp.pt")
+    
+    with open(nn_dir + "mlp.txt", "a") as writer:
+        writer.write(str(epoch) + "," + str(avg_loss) 
+                     + "," + str(avg_vloss) + "," + str(get_lr(optimizer)) + "\n")
+    
+    scheduler.step()
+

data/x_pointwise.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:277519100cb00cb7e059fd620358ac67ffe85a8cbb1daf0565f028ed68a76000
+size 1138964

data/y_pointwise.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50b39995e959275be47e658fd483bc169137705bb3ca8859272150767f0ca867
+size 284930