from heading import *
import torch
class CONV_net(torch.nn.Module):
def init(self, input_channels = 3, output_size = 10):
super().init()
self.conv1 = torch.nn.Sequential(
torch.nn.Conv2d(input_channels,16,3,2,1),
torch.nn.MaxPool2d(2),
torch.nn.ReLU(),
torch.nn.BatchNorm2d(16)
)
self.conv2 = torch.nn.Sequential(
torch.nn.Conv2d(16,32,3,1,1),
torch.nn.MaxPool2d(2),
torch.nn.ReLU(),
torch.nn.BatchNorm2d(32)
)
self.fc1 = torch.nn.Sequential(
torch.nn.Linear(288,60),
torch.nn.Dropout(0.5),
torch.nn.ReLU()
)
self.fc2 = torch.nn.Sequential(
torch.nn.Linear(60,20),
torch.nn.Dropout(0.5),
torch.nn.ReLU()
)
self.fc3 = torch.nn.Linear(20, output_size)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = x.view(-1,288)
x = self.fc1(x)
x = self.fc2(x)
x = self.fc3(x)
return x