20 lines
580 B
Python
20 lines
580 B
Python
import torch.nn as nn
|
|
|
|
|
|
class NeuralNet(nn.Module):
|
|
def __init__(self, input_size, hidden_size, num_classes):
|
|
super(NeuralNet, self).__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()
|
|
|
|
def forward(self, x):
|
|
out = self.l1(x)
|
|
out = self.relu(out)
|
|
out = self.l2(out)
|
|
out = self.relu(out)
|
|
out = self.l3(out)
|
|
# no activation and no softmax at the end
|
|
return out
|