I’m trying to build a neural network includes an SDEs layer and full connected NN. The problem that it takes hours just to complete one epoch.
This is my code:
First the SDEs class
class SDE(nn.Module):
def __init__(self):
super().__init__()
self.theta = nn.Parameter(torch.tensor(0.1), requires_grad=True) # Scalar parameter.
self.noise_type = "diagonal"
self.sde_type = "ito"
def f(self, t, y):
return torch.sin(t) + self.theta * y
def g(self, t, y):
return 0.3 * torch.sigmoid(torch.cos(t) * torch.exp(-y))
this is the pyTorch model:
class SDENET__(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(SDENET_, self).__init__()
self.batch_size, self.state_size, self.t_size = 1, 1, 16
self.ys = torch.tensor(0.1)
self.fc = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
sde = SDE()
ts = torch.linspace(0, 1, self.t_size)
y0 = torch.full(size=(self.batch_size, self.state_size), fill_value=0.1)
with torch.no_grad():
self.ys = torchsde.sdeint(sde, y0, ts, method='euler')
ys = self.ys[:, -1, :]
out = self.fc(ys)
return out
the model is to train NN for multivariate timeseries forecating.