72 lines
2.0 KiB
Python
72 lines
2.0 KiB
Python
import mlflow
|
|
import numpy as np
|
|
from tinygrad import Device,Tensor,nn,TinyJit
|
|
import matplotlib.pyplot as plt
|
|
import time
|
|
import show
|
|
from model import gen
|
|
|
|
BATCH_SIZE = 16
|
|
EPOCHS = 100
|
|
LEARNING_RATE = 1e-5
|
|
print(Device.DEFAULT)
|
|
mdl = gen()
|
|
opt = nn.optim.AdamW(nn.state.get_parameters(mdl), lr=LEARNING_RATE)
|
|
volume = 0.1
|
|
|
|
def spec_loss(pred, target, eps=1e-6):
|
|
# spectral convergence
|
|
sc = ((target - pred).square().sum()) ** 0.5 / ((target.square().sum()) ** 0.5 + eps)
|
|
# log magnitude difference
|
|
log_mag = ((target.abs() + eps).log() - (pred.abs() + eps).log()).abs().mean()
|
|
return sc + log_mag
|
|
|
|
|
|
@TinyJit
|
|
def step_gen(x):
|
|
Tensor.training = True
|
|
noise = Tensor.rand_like(x).tanh()
|
|
y = x+(noise*volume)
|
|
y = y.clamp(0,1)
|
|
loss = spec_loss(mdl(y),x)
|
|
opt.zero_grad()
|
|
loss.backward()
|
|
opt.step()
|
|
return loss.numpy()
|
|
|
|
print("loading")
|
|
x = np.load("data.npz")["arr_0"]
|
|
#x= x[0:64]
|
|
run_name = f"tinygrad_autoencoder_{int(time.time())}"
|
|
mlflow.set_tracking_uri("http://127.0.0.1:5000")
|
|
mlflow.start_run()
|
|
mlflow.log_params({"batch_size": BATCH_SIZE, "epochs": EPOCHS, "lr": LEARNING_RATE, "data size":len(x)})
|
|
|
|
show.logSpec(Tensor(x[0:1]).numpy()[0][0],"default")
|
|
|
|
print("training")
|
|
pl = 0
|
|
eshape = (BATCH_SIZE, 1, 128, 431)
|
|
for epoch in range(0,EPOCHS):
|
|
print(f"\n--- Starting Epoch {epoch} ---\n")
|
|
loss=0
|
|
for i in range(0,len(x),BATCH_SIZE):
|
|
tx=Tensor(x[i:i+BATCH_SIZE])
|
|
if(tx.shape != eshape):
|
|
continue
|
|
loss += step_gen(tx)
|
|
|
|
loss /= (len(x)/BATCH_SIZE)
|
|
if epoch%5==0:
|
|
noise = Tensor.rand_like(Tensor(x[0:1])).tanh()
|
|
y = Tensor(x[0:1]) + (noise*volume)
|
|
show.logSpec(mdl(y).numpy()[0][0],epoch)
|
|
if(pl - loss < 0.03 and epoch > 25):
|
|
show.logSpec(y.numpy()[0][0],f"volume_{volume}")
|
|
volume *= 2
|
|
pl = loss
|
|
|
|
mlflow.log_metric("volume", volume, step=epoch)
|
|
mlflow.log_metric("loss", loss, step=epoch)
|
|
print(f"loss of {loss}")
|