Rewrote dataset to be able to include new features

This commit is contained in:
Victor Mylle
2023-11-08 23:17:47 +00:00
parent 56c763a6f4
commit 2f48363292
10 changed files with 311218 additions and 118 deletions

View File

@@ -1,7 +1,7 @@
import torch
class NonLinearRegression(torch.nn.Module):
def __init__(self, inputSize, output_size, hiddenSize=128, numLayers=2):
def __init__(self, inputSize, output_size, hiddenSize=128, numLayers=2, dropout=0.0):
super(NonLinearRegression, self).__init__()
self.inputSize = inputSize
self.output_size = output_size
@@ -9,11 +9,15 @@ class NonLinearRegression(torch.nn.Module):
self.hiddenSize = hiddenSize
self.numLayers = numLayers
self.dropout = dropout
# add linear layers with relu
self.layers = torch.nn.ModuleList()
self.layers.append(torch.nn.Linear(inputSize, hiddenSize))
self.layers.append(torch.nn.Dropout(dropout))
for _ in range(numLayers - 2):
self.layers.append(torch.nn.Linear(hiddenSize, hiddenSize))
self.layers.append(torch.nn.Dropout(dropout))
self.layers.append(torch.nn.Linear(hiddenSize, output_size))
self.relu = torch.nn.ReLU()