|
@ -0,0 +1,16 @@ |
|
|
|
|
|
import torch.nn as nn |
|
|
|
|
|
import torch.nn.functional as F |
|
|
|
|
|
|
|
|
|
|
|
# CNN architecture definition |
|
|
|
|
|
class Net(nn.Module): |
|
|
|
|
|
def __init__(self): |
|
|
|
|
|
super(Net, self).__init__() |
|
|
|
|
|
# convolutional layer |
|
|
|
|
|
self.conv1 = nn.Conv2d(3, 16, 3, padding=1) |
|
|
|
|
|
# max pooling layer |
|
|
|
|
|
self.pool = nn.MaxPool2d(2, 2) |
|
|
|
|
|
|
|
|
|
|
|
def forward(self, x): |
|
|
|
|
|
# add sequence of convolutional and max pooling layers |
|
|
|
|
|
x = self.pool(F.relu(self.conv1(x))) |
|
|
|
|
|
return x |