29 lines
1.1 KiB
Python
29 lines
1.1 KiB
Python
|
import torch.nn as nn
|
|||
|
|
import torch.nn.functional as F
|
|||
|
|
|
|||
|
|
class NeuralNetwork(nn.Module):
|
|||
|
|
def __init__(self):
|
|||
|
|
super(NeuralNetwork, self).__init__()
|
|||
|
|
|
|||
|
|
# Warstwy konwolucyjnej sieci neuronowej, filtr 5×5, 3 kanały dla RGB
|
|||
|
|
self.convolutional_nn_1 = nn.Conv2d(3, 6, 5)
|
|||
|
|
self.convolutional_nn_2 = nn.Conv2d(6, 16, 5)
|
|||
|
|
|
|||
|
|
# Wyciaganie "najwazniejszej" informacji z obrazu
|
|||
|
|
self.pool = nn.MaxPool2d(2, 2)
|
|||
|
|
|
|||
|
|
self.full_connection_layer_1 = nn.Linear(16 * 71 * 71, 120)
|
|||
|
|
self.full_connection_layer_2 = nn.Linear(120, 84)
|
|||
|
|
self.full_connection_layer_3 = nn.Linear(84, 4)
|
|||
|
|
|
|||
|
|
# Forward określa przepływ inputu przez warstwy
|
|||
|
|
def forward(self, x):
|
|||
|
|
x = self.pool(F.relu(self.convolutional_nn_1(x)))
|
|||
|
|
x = self.pool(F.relu(self.convolutional_nn_2(x)))
|
|||
|
|
|
|||
|
|
# 16 kanałów o rozmiarach 71x71
|
|||
|
|
x = x.view(x.size(0), 16 * 71 * 71)
|
|||
|
|
x = F.relu(self.full_connection_layer_1(x))
|
|||
|
|
x = F.relu(self.full_connection_layer_2(x))
|
|||
|
|
x = self.full_connection_layer_3(x)
|
|||
|
|
return x
|