tony/tensorneat-mend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

complete fully stateful!

Browse Source 
use black to format all files!
This commit is contained in:
wls2002
2024-05-26 18:08:43 +08:00
parent cf69b916af
commit 18c3d44c79
41 changed files with 620 additions and 495 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
tensorneat/problem/func_fit/func_fit.py
View File

@@ -8,42 +8,44 @@ from .. import BaseProblem
class FuncFit(BaseProblem):
jitable = True
def __init__(self,
error_method: str = 'mse'
):
def __init__(self, error_method: str = "mse"):
super().__init__()
assert error_method in {'mse', 'rmse', 'mae', 'mape'}
assert error_method in {"mse", "rmse", "mae", "mape"}
self.error_method = error_method
def setup(self, state: State = State()):
return state
def evaluate(self, randkey, state, act_func, params):
def evaluate(self, state, randkey, act_func, params):
state, predict = jax.vmap(act_func, in_axes=(None, 0, None), out_axes=(None, 0))(state, self.inputs, params)
predict = jax.vmap(act_func, in_axes=(None, 0, None))(
state, self.inputs, params
)
if self.error_method == 'mse':
if self.error_method == "mse":
loss = jnp.mean((predict - self.targets) ** 2)
elif self.error_method == 'rmse':
elif self.error_method == "rmse":
loss = jnp.sqrt(jnp.mean((predict - self.targets) ** 2))
elif self.error_method == 'mae':
elif self.error_method == "mae":
loss = jnp.mean(jnp.abs(predict - self.targets))
elif self.error_method == 'mape':
elif self.error_method == "mape":
loss = jnp.mean(jnp.abs((predict - self.targets) / self.targets))
else:
raise NotImplementedError
return state, -loss
return -loss
def show(self, randkey, state, act_func, params, *args, **kwargs):
state, predict = jax.vmap(act_func, in_axes=(None, 0, None), out_axes=(None, 0))(state, self.inputs, params)
def show(self, state, randkey, act_func, params, *args, **kwargs):
predict = jax.vmap(act_func, in_axes=(None, 0, None))(
state, self.inputs, params
)
inputs, target, predict = jax.device_get([self.inputs, self.targets, predict])
state, loss = self.evaluate(randkey, state, act_func, params)
loss = self.evaluate(state, randkey, act_func, params)
loss = -loss
msg = ""