File size: 670 Bytes
42e3a78 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
import torch.nn as nn
class NeuralNet(nn.Module):
def __init__(self, input_size, hidden_size, num_classes):
super().__init__()
self.l1 = nn.Linear(input_size, hidden_size)
self.l2 = nn.Linear(hidden_size, hidden_size)
self.l3 = nn.Linear(hidden_size, num_classes)
self.relu = nn.ReLU()
self.dropout = nn.Dropout(p=0.5)
def forward(self, x):
out = self.l1(x)
out = self.relu(out)
out = self.dropout(out)
out = self.l2(out)
out = self.relu(out)
out = self.dropout(out)
out = self.l3(out)
# no activation and no softmax at the end
return out
|