27 lines
683 B
Python
27 lines
683 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class NeuralNetwork(nn.Module):
|
|
"""
|
|
Neural network model for classification problem.
|
|
"""
|
|
|
|
def __init__(self, input_size, hidden_size):
|
|
super(NeuralNetwork, self).__init__()
|
|
|
|
self.fc1 = nn.Linear(input_size, hidden_size)
|
|
self.fc2 = nn.Linear(hidden_size, hidden_size // 2)
|
|
self.fc3 = nn.Linear(hidden_size // 2, 1)
|
|
|
|
self.relu = nn.ReLU()
|
|
self.sigmoid = nn.Sigmoid()
|
|
|
|
def forward(self, x):
|
|
out = self.fc1(x)
|
|
out = self.relu(out)
|
|
out = self.fc2(out)
|
|
out = self.relu(out)
|
|
out = self.fc3(out)
|
|
out = self.sigmoid(out)
|
|
return out |