54 lines
2.0 KiB
Python
54 lines
2.0 KiB
Python
import torch
|
|
import torch.nn
|
|
|
|
|
|
class MyCNN(torch.nn.Module):
|
|
def __init__(self,
|
|
input_channels: int):
|
|
super().__init__()
|
|
|
|
self.layers = torch.nn.Sequential(
|
|
torch.nn.Conv2d(in_channels=input_channels, out_channels=32, kernel_size=3, padding='same', bias=False),
|
|
torch.nn.BatchNorm2d(32),
|
|
torch.nn.ReLU(),
|
|
torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding='same', bias=False),
|
|
torch.nn.BatchNorm2d(32),
|
|
torch.nn.ReLU(),
|
|
torch.nn.MaxPool2d(kernel_size=3, padding=1),
|
|
|
|
torch.nn.Dropout2d(0.1),
|
|
torch.nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding='same', bias=False),
|
|
torch.nn.BatchNorm2d(64),
|
|
torch.nn.ReLU(),
|
|
torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding='same', bias=False),
|
|
torch.nn.BatchNorm2d(64),
|
|
torch.nn.ReLU(),
|
|
torch.nn.MaxPool2d(kernel_size=3, padding=1),
|
|
|
|
torch.nn.Dropout2d(0.1),
|
|
torch.nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding='same', bias=False),
|
|
torch.nn.BatchNorm2d(128),
|
|
torch.nn.ReLU(),
|
|
torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, padding='same', bias=False),
|
|
torch.nn.BatchNorm2d(128),
|
|
torch.nn.ReLU(),
|
|
torch.nn.MaxPool2d(kernel_size=3, padding=1),
|
|
|
|
torch.nn.Flatten(),
|
|
torch.nn.Dropout(0.25),
|
|
torch.nn.Linear(in_features=2048, out_features=1024),
|
|
torch.nn.ReLU(),
|
|
torch.nn.Linear(in_features=1024, out_features=512),
|
|
torch.nn.ReLU(),
|
|
torch.nn.Linear(in_features=512, out_features=20, bias=False)
|
|
)
|
|
|
|
def forward(self, input_images: torch.Tensor) -> torch.Tensor:
|
|
return self.layers(input_images)
|
|
|
|
def __repr__(self):
|
|
return str(self.layers)
|
|
|
|
|
|
model = MyCNN(input_channels=1)
|