from utils import * class Model_Vehicle_CLS_2_1(nn.Module): def __init__(self, num_classes=3): super(Model_Vehicle_CLS_2_1, self).__init__() self.conv1 = nn.Sequential( nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1, dilation=1, bias=False), nn.BatchNorm2d(16), nn.ReLU(inplace=True), nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1, dilation=1, bias=False), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1, dilation=1, bias=False), nn.BatchNorm2d(64), nn.ReLU(inplace=True), ) self.conv2 = nn.Sequential( nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, dilation=1, bias=False), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1, dilation=1, bias=False), nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1, dilation=1, bias=False), nn.BatchNorm2d(512), nn.ReLU(inplace=True), ) self.fc = nn.Linear(in_features=512, out_features=num_classes) def forward(self, x): x = self.conv1(x) x = self.conv2(x) x = F.avg_pool2d(x, 32) x = x.view(x.size(0), -1) x = self.fc(x) return x class Model_Vehicle_CLS_2_2(nn.Module): def __init__(self, num_classes=3): super(Model_Vehicle_CLS_2_2, self).__init__() self.conv1 = nn.Sequential( nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1, dilation=1, bias=False), nn.BatchNorm2d(16), nn.ReLU(inplace=True), nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=2, dilation=2, bias=False), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=5, dilation=5, bias=False), nn.BatchNorm2d(64), nn.ReLU(inplace=True), ) self.conv2 = nn.Sequential( nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, dilation=1, bias=False), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=2, dilation=2, bias=False), nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=5, dilation=5, bias=False), nn.BatchNorm2d(512), nn.ReLU(inplace=True), ) self.fc = nn.Linear(in_features=512, out_features=num_classes) def forward(self, x): x = self.conv1(x) x = self.conv2(x) x = F.avg_pool2d(x, 32) x = x.view(x.size(0), -1) x = self.fc(x) return x class Model_Vehicle_CLS_2_3(nn.Module): def __init__(self, num_classes=3): super(Model_Vehicle_CLS_2_3, self).__init__() self.conv1 = nn.Sequential( nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1, dilation=1, bias=False), nn.BatchNorm2d(16), nn.ReLU(inplace=True), nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=3, dilation=3, bias=False), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=5, dilation=5, bias=False), nn.BatchNorm2d(64), nn.ReLU(inplace=True), ) self.conv2 = nn.Sequential( nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, dilation=1, bias=False), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=3, dilation=3, bias=False), nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=5, dilation=5, bias=False), nn.BatchNorm2d(512), nn.ReLU(inplace=True), ) self.fc = nn.Linear(in_features=512, out_features=num_classes) def forward(self, x): x = self.conv1(x) x = self.conv2(x) x = F.avg_pool2d(x, 32) x = x.view(x.size(0), -1) x = self.fc(x) return x class Model_Vehicle_CLS_2_4(nn.Module): def __init__(self, num_classes=3): super(Model_Vehicle_CLS_2_4, self).__init__() self.conv1 = nn.Sequential( nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1, dilation=1, bias=False), nn.BatchNorm2d(16), nn.ReLU(inplace=True), nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=3, dilation=3, bias=False), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=7, dilation=7, bias=False), nn.BatchNorm2d(64), nn.ReLU(inplace=True), ) self.conv2 = nn.Sequential( nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, dilation=1, bias=False), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=3, dilation=3, bias=False), nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=7, dilation=7, bias=False), nn.BatchNorm2d(512), nn.ReLU(inplace=True), ) self.fc = nn.Linear(in_features=512, out_features=num_classes) def forward(self, x): x = self.conv1(x) x = self.conv2(x) x = F.avg_pool2d(x, 32) x = x.view(x.size(0), -1) x = self.fc(x) return x if __name__ == "__main__": num_epochs = 61 learning_rate = 1e-4 batch_size = 256 dilations = ["[[1, 1, 1], [1, 1, 1]]", "[[1, 2, 5], [1, 2, 5]]", "[[1, 3, 5], [1, 3, 5]]", "[[1, 3, 7], [1, 3, 7]]"] models = [ Model_Vehicle_CLS_2_1, Model_Vehicle_CLS_2_2, Model_Vehicle_CLS_2_3, Model_Vehicle_CLS_2_4, ] for i in range(4): model = models[i]() print("Dilation: " + dilations[i]) train_loss, test_acc = train_Vehicle_CLS(model=model, learning_rate=learning_rate, batch_size=batch_size, num_epochs=num_epochs) print()