Skip to content

Commit

Permalink
Refactor code for speed and clarity
Browse files Browse the repository at this point in the history
  • Loading branch information
@glenn-jocher
glenn-jocher committed Jun 9, 2024
1 parent cbd9b9a commit 938e3f2
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 35 deletions.
45 changes: 19 additions & 26 deletions gcp/wave_pytorch_gcp.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@

import scipy.io
import torch

from utils import *

# set printoptions
Expand All @@ -24,15 +23,15 @@ def runexample(H, model, str, lr=0.001, amsgrad=False):
data = "wavedata25ns.mat"

cuda = torch.cuda.is_available()
os.makedirs(pathr + "models", exist_ok=True)
os.makedirs(f"{pathr}models", exist_ok=True)
name = (data[:-4] + "%s%glr%s" % (H[:], lr, str)).replace(", ", ".").replace("[", "_").replace("]", "_")

tica = time.time()
device = torch.device("cuda:0" if cuda else "cpu")
print("Running %s on %s\n%s" % (name, device.type, torch.cuda.get_device_properties(0) if cuda else ""))

if not os.path.isfile(pathd + data):
os.system("wget -P data/ https://storage.googleapis.com/ultralytics/" + data)
os.system(f"wget -P data/ https://storage.googleapis.com/ultralytics/{data}")
mat = scipy.io.loadmat(pathd + data)
x = mat["inputs"] # inputs (nx512) [waveform1 waveform2]
y = mat["outputs"][:, 1:2] # outputs (nx4) [position(mm), time(ns), PE, E(MeV)]
Expand Down Expand Up @@ -146,7 +145,7 @@ class LinearAct(torch.nn.Module):
def __init__(self, nx, ny):
super(LinearAct, self).__init__()
self.Linear1 = torch.nn.Linear(nx, ny)
self.act = eval("torch.nn." + a + "()")
self.act = eval(f"torch.nn.{a}()")

def forward(self, x):
return self.act(self.Linear1(x))
Expand All @@ -161,9 +160,9 @@ def __init__(self, n): # n = [512, 108, 23, 5, 1]
def forward(self, x):
return self.fc2(self.fc1(self.fc0(x)))

for i in range(10):
tsy.append(runexample(H, model=WAVE(H), str=("." + a)))
scipy.io.savemat(pathr + "TS.sigmoid" + ".mat", dict(tsv=tsv, tsy=np.array(tsy)))
for _ in range(10):
tsy.append(runexample(H, model=WAVE(H), str=f".{a}"))
scipy.io.savemat(f"{pathr}TS.sigmoid.mat", dict(tsv=tsv, tsy=np.array(tsy)))


def tsnoact(): # TS activation function
Expand All @@ -184,27 +183,25 @@ def __init__(self, n): # n = [512, 108, 23, 5, 1]
def forward(self, x):
return self.fc2(self.fc1(self.fc0(x)))

for i in range(10):
tsy.append(runexample(H, model=WAVE(H), str=("." + a)))
scipy.io.savemat(pathr + "TS.noact" + ".mat", dict(tsv=tsv, tsy=np.array(tsy)))
for _ in range(10):
tsy.append(runexample(H, model=WAVE(H), str=f".{a}"))
scipy.io.savemat(f"{pathr}TS.noact.mat", dict(tsv=tsv, tsy=np.array(tsy)))


def tslr(): # TS learning rate
tsv = np.logspace(-5, -2, 13)
tsy = []
for a in tsv:
for i in range(10):
tsy.append(runexample(H, model=WAVE(H), str=("." + "Tanh"), lr=a))
scipy.io.savemat(pathr + "TS.lr" + ".mat", dict(tsv=tsv, tsy=np.array(tsy)))
tsy.extend(runexample(H, model=WAVE(H), str=("." + "Tanh"), lr=a) for _ in range(10))
scipy.io.savemat(f"{pathr}TS.lr.mat", dict(tsv=tsv, tsy=np.array(tsy)))


def tsams(): # TS AMSgrad
tsv = [False, True]
tsy = []
for a in tsv:
for i in range(3):
tsy.append(runexample(H, model=WAVE(H), str=(".TanhAMS" + str(a)), amsgrad=a))
scipy.io.savemat(pathr + "TS.AMSgrad" + ".mat", dict(tsv=tsv, tsy=np.array(tsy)))
tsy.extend(runexample(H, model=WAVE(H), str=f".TanhAMS{str(a)}", amsgrad=a) for _ in range(3))
scipy.io.savemat(f"{pathr}TS.AMSgrad.mat", dict(tsv=tsv, tsy=np.array(tsy)))


def tsshape(): # TS network shape
Expand All @@ -223,8 +220,6 @@ def tsshape(): # TS network shape
# tsv = ['Tanh', 'LogSigmoid', 'Softsign', 'ELU']
# tsv = np.logspace(-4, -2, 11)
tsv = [[512, 23, 1], [512, 64, 8, 1], [512, 108, 23, 5, 1], [512, 147, 42, 12, 3, 1], [512, 181, 64, 23, 8, 3, 1]]
tsy = []

H = tsv[0]

class WAVE(torch.nn.Module):
Expand All @@ -236,9 +231,7 @@ def __init__(self, n): # n = [512, 108, 23, 5, 1]
def forward(self, x):
return self.fc1(self.fc0(x))

for i in range(10):
tsy.append(runexample(H, model=WAVE(H), str=("." + "Tanh")))

tsy = [runexample(H, model=WAVE(H), str=("." + "Tanh")) for _ in range(10)]
H = tsv[1]

class WAVE(torch.nn.Module):
Expand All @@ -251,7 +244,7 @@ def __init__(self, n): # n = [512, 108, 23, 5, 1]
def forward(self, x):
return self.fc2(self.fc1(self.fc0(x)))

for i in range(10):
for _ in range(10):
tsy.append(runexample(H, model=WAVE(H), str=("." + "Tanh")))

H = tsv[2]
Expand All @@ -267,7 +260,7 @@ def __init__(self, n): # n = [512, 108, 23, 5, 1]
def forward(self, x):
return self.fc3(self.fc2(self.fc1(self.fc0(x))))

for i in range(10):
for _ in range(10):
tsy.append(runexample(H, model=WAVE(H), str=("." + "Tanh")))

H = tsv[3]
Expand All @@ -284,7 +277,7 @@ def __init__(self, n): # n = [512, 108, 23, 5, 1]
def forward(self, x):
return self.fc4(self.fc3(self.fc2(self.fc1(self.fc0(x)))))

for i in range(10):
for _ in range(10):
tsy.append(runexample(H, model=WAVE(H), str=("." + "Tanh")))

H = tsv[4]
Expand All @@ -302,9 +295,9 @@ def __init__(self, n): # n = [512, 108, 23, 5, 1]
def forward(self, x):
return self.fc5(self.fc4(self.fc3(self.fc2(self.fc1(self.fc0(x))))))

for i in range(10):
for _ in range(10):
tsy.append(runexample(H, model=WAVE(H), str=("." + "Tanh")))
scipy.io.savemat(pathr + "TS.shape" + ".mat", dict(tsv=tsv, tsy=np.array(tsy)))
scipy.io.savemat(f"{pathr}TS.shape.mat", dict(tsv=tsv, tsy=np.array(tsy)))


if __name__ == "__main__":
Expand Down
7 changes: 3 additions & 4 deletions train.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@

import scipy.io
import torch.nn as nn

from utils.torch_utils import *
from utils.utils import *

Expand All @@ -21,14 +20,14 @@ def train(H, model, str, lr=0.001):
data = "wavedata25ns.mat"

cuda = torch.cuda.is_available()
os.makedirs(pathr + "models", exist_ok=True)
os.makedirs(f"{pathr}models", exist_ok=True)
name = (data[:-4] + "%s%glr%s" % (H[:], lr, str)).replace(", ", ".").replace("[", "_").replace("]", "_")
print("Running " + name)
print(f"Running {name}")

device = select_device()

if not os.path.isfile(pathd + data):
os.system("wget -P data/ https://storage.googleapis.com/ultralytics/" + data)
os.system(f"wget -P data/ https://storage.googleapis.com/ultralytics/{data}")
mat = scipy.io.loadmat(pathd + data)
x = mat["inputs"][:] # inputs (nx512) [waveform1 waveform2]
y = mat["outputs"][:, 0:2] # outputs (nx4) [position(mm), time(ns), PE, E(MeV)]
Expand Down
9 changes: 4 additions & 5 deletions train_tf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
import tensorflow as tf
import tensorflow.contrib.eager as tfe
from plotly.offline import plot

from utils.utils import *

tf.enable_eager_execution()
Expand All @@ -24,15 +23,15 @@ def runexample(H, model, str):
cuda = tf.test.is_gpu_available()
tf.set_random_seed(1)
path = "data/"
os.makedirs(path + "models", exist_ok=True)
os.makedirs(f"{path}models", exist_ok=True)
name = (data[:-4] + "%s%glr%geps%s" % (H[:], lr, eps, str)).replace(", ", "_").replace("[", "_").replace("]", "_")

tica = time.time()
device = "/gpu:0" if cuda else "/cpu:0"
print("Running %s on %s" % (name, device))
print(f"Running {name} on {device}")

if not os.path.isfile(path + data):
os.system("wget -P data/ https://storage.googleapis.com/ultralytics/" + data)
os.system(f"wget -P data/ https://storage.googleapis.com/ultralytics/{data}")
mat = scipy.io.loadmat(path + data)
x = mat["inputs"] # inputs (nx512) [waveform1 waveform2]
y = mat["outputs"][:, 0:2] # outputs (nx4) [position(mm), time(ns), PE, E(MeV)]
Expand Down Expand Up @@ -128,4 +127,4 @@ def criteria(y_pred, y): # MSE
if __name__ == "__main__":
H = [128, 32, 8]
for i in range(1):
runexample(H, None, "." + str(i))
runexample(H, None, f".{str(i)}")

0 comments on commit 938e3f2

Please sign in to comment.