I want to define a custom loss function, with a little modification to RMSE
I want to calculate the rmse of pred head and tail
here is the code:
from util import *
import torch.optim as opt
import torch.nn as nn
import torch
def getxy(seed=100):
np.random.seed(100)
x = np.random.rand(100, 2)
w = np.array([3, 7])
b = 5
y = (x * w).sum(axis=1) + b ### remember, keep the shape same with network output!!!
# y = 3 * x1 + 7* x2 + 5
return x, y.reshape(-1,1)
class MLP(nn.Module):
def __init__(self):
super(MLP, self).__init__()
self.l1 = nn.Linear(2, 1, bias=True)
def forward(self, x):
return self.l1(torch.tensor(x).double())
class myloss(nn.Module):
def __init__(self):
super().__init__()
#TODO: HOW CAN I MAKE THIS WORK??
def forward(self, x, y):
xarray = x.detach().numpy()
mask = (xarray > np.quantile(xarray, 0.9, axis=0)) + (xarray < np.quantile(xarray, 0.1, axis=0))
xarray = np.where(mask, xarray, y)
l = mask.sum()
#return ((xarray - y) ** 2).mean(axis=1)
#return torch.tensor(torch.sum(torch.square(torch.tensor(xarray) - torch.tensor(y))) / l)
return torch.mean(torch.square(torch.tensor(xarray) - torch.tensor(y)))
#def backward(grad_output):
#return grad_output
mlp = MLP().double()
loss = myloss()
o = opt.Adam(mlp.parameters(), lr=0.1)
#o = opt.SGD(mlp.parameters(), lr=0.05)
def train(n):
x, y = getxy()
for i in range(n):
pred = mlp(x)
o.zero_grad()
los = loss(pred, y)
los.backward()
o.step()
print('%dth: loss = %.3f'%(i, los.item()), mlp.l1.weight.detach().numpy(), mlp.l1.bias.detach().numpy())
if __name__ == '__main__':
train(10000)
Can anyone help on this? or provide a demo which can use numpy in custom loss function freely?