Rewrote dataset to be able to include new features
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user