52 lines
1.5 KiB
Python
52 lines
1.5 KiB
Python
|
import tensorflow as tf
|
||
|
import numpy as np
|
||
|
from matplotlib import pyplot as plt
|
||
|
import itertools as it
|
||
|
|
||
|
from sklearn.metrics import r2_score
|
||
|
|
||
|
from math import sin
|
||
|
|
||
|
def f(x, y): return (1 + sin(x**2 + 5*y)) / 2
|
||
|
|
||
|
def train_generic(model, marker):
|
||
|
model.compile(optimizer = "adam", loss = "mse")
|
||
|
|
||
|
X = np.linspace(0, 10, 300)
|
||
|
Y = np.linspace(0, 10, 300)
|
||
|
|
||
|
ins = np.array(list(it.product(X, Y)))
|
||
|
ous = np.array(list(f(*i) for i in ins))
|
||
|
|
||
|
result = model.fit(ins, ous, epochs = 200, batch_size = 2048)
|
||
|
|
||
|
model.save_weights(f"save-{marker}.weights.h5")
|
||
|
|
||
|
def verify_generic(model, marker):
|
||
|
model.compile(optimizer = "adam", loss = "mse")
|
||
|
model.load_weights(f"save-{marker}.weights.h5")
|
||
|
|
||
|
X = np.linspace(0, 10, 50)
|
||
|
Y = np.linspace(0, 10, 50)
|
||
|
|
||
|
ins = np.array(list(it.product(X, Y)))
|
||
|
ous = np.array(list(f(*i) for i in ins))
|
||
|
|
||
|
preds = model.predict(ins)
|
||
|
|
||
|
print(f"Model {marker} has {r2_score(ous, preds)} r2 score")
|
||
|
|
||
|
preds_flat = [i[0] for i in preds]
|
||
|
|
||
|
px = np.array([X] * len(X))
|
||
|
py = np.array([[x] * (len(X)) for x in X])
|
||
|
pz1 = np.array([ous[i*len(X):(i+1)*len(X)] for i, _ in enumerate(X)])
|
||
|
pz2 = np.array([preds_flat[i*len(X):(i+1)*len(X)] for i, _ in enumerate(X)])
|
||
|
|
||
|
#print([ous[i*len(X):(i+1)*len(X)] for i, _ in enumerate(X)])
|
||
|
|
||
|
p = plt.figure().add_subplot(projection='3d')
|
||
|
p.plot_surface(px, py, pz1, edgecolor = "lime", alpha = 0.1)
|
||
|
p.plot_surface(px, py, pz2, edgecolor = "blue", alpha = 0.3)
|
||
|
plt.show()
|