84 lines
2.2 KiB
Python
84 lines
2.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
VGG_CONFIG = {
|
|
"VGG16": [64, 64, "M",
|
|
128, 128, "M",
|
|
256, 256, 256, "M",
|
|
512, 512, 512, "M",
|
|
512, 512, 512, "M"],
|
|
"VGG19": [64, 64, "M",
|
|
128, 128, "M",
|
|
256, 256, 256, 256, "M",
|
|
512, 512, 512, 512, "M",
|
|
512, 512, 512, 512, "M"]
|
|
}
|
|
|
|
|
|
class VGGCustom(nn.Module):
|
|
def __init__(self,
|
|
architecture:str="VGG16",
|
|
num_channels:int = 3,
|
|
num_classes:int=10,
|
|
batch_normalization=False,
|
|
dropout:float=0.5) -> None:
|
|
|
|
super().__init__()
|
|
|
|
if architecture not in VGG_CONFIG:
|
|
raise ValueError(f"Architecture {architecture} is invalid. ",
|
|
f"Available architectures are :{list(VGG_CONFIG.keys())}")
|
|
|
|
|
|
self.batch_normalization = batch_normalization
|
|
|
|
self.features = self._make_layers(
|
|
VGG_CONFIG[architecture],
|
|
in_channels=num_channels
|
|
)
|
|
|
|
self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
|
|
|
|
self.classifier = nn.Sequential(
|
|
nn.Linear(512*7*7, 4096),
|
|
nn.ReLU(True),
|
|
nn.Dropout(p=dropout),
|
|
|
|
nn.Linear(4096, 4096),
|
|
nn.ReLU(True),
|
|
nn.Dropout(p=dropout),
|
|
|
|
nn.Linear(4096, num_classes)
|
|
)
|
|
|
|
def _make_layers(self, cfg, in_channels):
|
|
layers = []
|
|
|
|
for layer in cfg:
|
|
if layer == 'M':
|
|
layers.append(nn.MaxPool2d(kernel_size=2, stride=2))
|
|
else:
|
|
|
|
layers.append(nn.Conv2d(in_channels=in_channels,
|
|
out_channels=layer,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1))
|
|
|
|
if self.batch_normalization:
|
|
layers.append(nn.BatchNorm2d(layer))
|
|
|
|
layers.append(nn.ReLU(True))
|
|
|
|
in_channels = layer
|
|
|
|
return nn.Sequential(*layers)
|
|
|
|
def forward(self, image):
|
|
x = self.features(image)
|
|
x = self.avgpool(x)
|
|
x = torch.flatten(x, 1)
|
|
x = self.classifier(x)
|
|
|
|
return x |