forked from s464965/WMICraft
23 lines
758 B
Python
23 lines
758 B
Python
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
import torch.nn.functional as F
|
||
|
|
|
||
|
|
|
||
|
|
class NeuralNetwork(nn.Module):
|
||
|
|
def __init__(self, num_classes=4):
|
||
|
|
super(NeuralNetwork, self).__init__()
|
||
|
|
self.conv1 = nn.Conv2d(in_channels=3, out_channels=10, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||
|
|
self.pool = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2))
|
||
|
|
self.conv2 = nn.Conv2d(in_channels=10, out_channels=20, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||
|
|
self.fc1 = nn.Linear(20*9*9, num_classes)
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
x = F.relu(self.conv1(x))
|
||
|
|
x = self.pool(x)
|
||
|
|
x = F.relu(self.conv2(x))
|
||
|
|
x = self.pool(x)
|
||
|
|
x = x.reshape(x.shape[0], -1)
|
||
|
|
x = self.fc1(x)
|
||
|
|
|
||
|
|
return x
|