import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy
import torch.nn.functional as F
RESNET = ['resnet18','resnet34','resnet50','resnet101','resnet152']
def set_parameter_requires_grad(model, feature_extracting):
if feature_extracting:
for param in model.parameters():
param.requires_grad = False
def initialize_model(conf):
model_name = conf['model']['name']
feature_extract = conf['model']['feature_extract']
use_pretrained = conf['model']['use_pretrained']
print_model = conf['model']['print_model']
num_classes = len(conf['data']['classes'])
if model_name in RESNET:
model = getattr(models, model_name)(pretrained=use_pretrained)
set_parameter_requires_grad(model, feature_extract)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, num_classes)
elif model_name == 'efficientdet_d0':
model = EfficientClassification(num_classes)
else:
print("Invalid model name, exiting...")
exit()
if print_model: print(model)
model.name = model_name
return model
class EfficientClassification(nn.Module):
def __init__(self, num_classes):
super(EfficientClassification, self).__init__()
from effdet import create_model
self.effdet = create_model(model_name='efficientdet_d0')
self.effdet.box_net = nn.Identity()
self.effdet.class_net = nn.Identity()
self.resnet = models.resnet18(pretrained=True)
num_ftrs = self.resnet.fc.in_features
self.resnet.fc = nn.Linear(num_ftrs, num_classes)
self.deconv0 = nn.ConvTranspose2d(in_channels=64,
out_channels=16,
kernel_size=19,
stride=3,
padding=1,
dilation=2)
self.deconv1 = nn.ConvTranspose2d(in_channels=64,
out_channels=12,
kernel_size=9,
stride=7,
padding=1,
dilation=1)
self.deconv2 = nn.ConvTranspose2d(in_channels=64,
out_channels=8,
kernel_size=24,
stride=9,
padding=2,
dilation=4)
self.deconv3 = nn.ConvTranspose2d(in_channels=64,
out_channels=4,
kernel_size=28,
stride=9,
padding=1,
dilation=6)
self.deconv4 = nn.ConvTranspose2d(in_channels=64,
out_channels=2,
kernel_size=30,
stride=8,
padding=2,
dilation=7)
self.conv0 = nn.Conv2d(in_channels=42,
out_channels=16,
kernel_size=5,
padding=2)
self.conv1 = nn.Conv2d(in_channels=16,
out_channels=3,
kernel_size=3,
padding=1)
def forward(self, x):
# EffNet + BiFPN
fpn_out, _ = self.effdet(x)
# Convolution Transpose
out0 = self.deconv0(fpn_out[0])
out1 = self.deconv1(fpn_out[1])
out2 = self.deconv2(fpn_out[2])
out3 = self.deconv3(fpn_out[3])
out4 = self.deconv4(fpn_out[4])
deconv_out = torch.cat([out0,out1,out2,out3,out4], dim=1)
# Convolution
conv_out = self.conv1(self.conv0(deconv_out))
# Resnet18
out = self.resnet(conv_out)
return out
class EfficientClassification2(nn.Module):
def __init__(self, num_classes):
super(EfficientClassification2, self).__init__()
from effdet import create_model
self.effdet = create_model(model_name='efficientdet_d0')
self.effdet.box_net = nn.Identity()
self.effdet.class_net = nn.Identity()
# In features from FPN
fc_in_features = [64 * i*i for i in [64,32,16,8,4]]
mid = 64
self.fc0 = nn.Linear(fc_in_features[0], mid)
self.fc1 = nn.Linear(fc_in_features[1], mid)
self.fc2 = nn.Linear(fc_in_features[2], mid)
self.fc3 = nn.Linear(fc_in_features[3], mid)
self.fc4 = nn.Linear(fc_in_features[4], mid)
self.fc_out = nn.Linear(5 * mid, num_classes)
def forward(self, x):
fpn_out, _ = self.effdet(x)
fpn_out = list(map(lambda t: torch.flatten(t, start_dim=1), fpn_out))
out0 = self.fc0(fpn_out[0])
out1 = self.fc1(fpn_out[1])
out2 = self.fc2(fpn_out[2])
out3 = self.fc3(fpn_out[3])
out4 = self.fc4(fpn_out[4])
fc_outs = torch.cat([out0,out1,out2,out3,out4], dim=1)
out = self.fc_out(fc_outs)
return out
if __name__ == '__main__':
x = torch.randn(20, 3, 512, 512)
model = EfficientClassification(num_classes=2)
fpn_out = model(x)
print('FIN')
Datasets
Models
