import torch.nn as nn
class CNN(nn.Module):
	def __init__(self):
		super(CNN,self).__init__()
		self.convLayer = nn.Sequential(
			nn.Conv2d(in_channels = 3,out_channels = 32,kernel_size = 3,padding = 1),
			nn.BatchNorm2d(32),
			nn.ReLU(inplace = True),
			nn.Conv2d(in_channels = 32,out_channels = 64,kernel_size = 3,padding = 1),
			nn.ReLU(inplace = True),
			nn.MaxPool2d(kernel_size = 2,stride = 2),


			nn.Conv2d(in_channels = 64,out_channels = 128,kernel_size = 3,padding = 1),
			nn.BatchNorm2d(128),
			nn.ReLU(inplace = True),
			nn.Conv2d(in_channels = 128,out_channels = 128,kernel_size = 3,padding = 1),
			nn.ReLU(inplace = True),
			nn.MaxPool2d(kernel_size = 2,stride = 2),
			nn.Dropout2d(p = 0.05),

			nn.Conv2d(in_channels = 128,out_channels = 256,kernel_size = 3,padding = 1),
			nn.BatchNorm2d(256),
			nn.ReLU(inplace = True),
			nn.MaxPool2d(kernel_size = 2,stride = 2),
			nn.Conv2d(in_channels = 256,out_channels = 256,kernel_size = 3,padding = 1),
			nn.ReLU(inplace = True),
			nn.MaxPool2d(kernel_size = 2,stride = 2))

			
		self.fc_layer = nn.Sequential(
			nn.Dropout(p = 0.1),
			nn.Linear(9216,1024),
			nn.ReLU(inplace = True),
			nn.Linear(1024,512),
			nn.ReLU(inplace = True),
			nn.Dropout(p = 0.1),
			nn.Linear(512,4)
			)

	def forward(self,x):
		x = self.convLayer(x)
		x = x.view(x.size(0),-1)
		return self.fc_layer(x)