def __init__(self, num_classes=11, checkpoint_path=None):
super(EfficientDet5AP, self).__init__()
config = get_efficientdet_config('tf_efficientdet_d5_ap')
config.image_size = [512, 512]
config.norm_kwargs = dict(eps=.001, momentum=.01)
config.soft_nms = True
config.label_smoothing = 0.1
# config.legacy_focal = True
net = EfficientDet(config, pretrained_backbone=False)
if checkpoint_path == None:
checkpoint = torch.load(
'./effdet_model/tf_efficientdet_d5_ap-3673ae5d.pth')
net.load_state_dict(checkpoint)
net.reset_head(num_classes=num_classes)
net.class_net = HeadNet(config, num_outputs=config.num_classes)
self.model = DetBenchTrain(net, config)
else:
checkpoint = torch.load(checkpoint_path)
checkpoint2 = {
'.'.join(k.split('.')[2:]): v
for k, v in checkpoint.items()
}
del checkpoint2['boxes']
net.reset_head(num_classes=num_classes)
net.class_net = HeadNet(config, num_outputs=config.num_classes)
net.load_state_dict(checkpoint2)
self.model = DetBenchPredict(net)
def load_detect(self, checkpoint_path):
config = get_efficientdet_config('tf_efficientdet_d0')
config.image_size = [384, 384]
net = EfficientDet(config, pretrained_backbone=False)
net.reset_head(num_classes=1)
checkpoint = torch.load(checkpoint_path,
map_location=torch.device('cpu'))
net.load_state_dict(checkpoint['model_state_dict'])
del checkpoint
gc.collect()
self.net = DetBenchPredict(net)
self.net.eval()
self.net.cpu()
def get_net():
# config = get_efficientdet_config('tf_efficientdet_d4')
config = get_efficientdet_config('tf_efficientdet_d0')
print(config)
net = EfficientDet(config, pretrained_backbone=False)
count = 0
for param in net.parameters():
count += torch.prod(torch.tensor(param.shape))
print(count)
# checkpoint = torch.load('/content/drive/MyDrive/efficientdet-pytorch/efficientdet_d4-5b370b7a.pth')
# checkpoint = torch.load('/content/drive/MyDrive/efficientdet-pytorch/efficientdet_d0-d92fd44f.pth')
# net.load_state_dict(checkpoint)
state_dict = load_state_dict_from_url(config.url,
progress=False,
map_location='cpu')
net.load_state_dict(state_dict, strict=True)
net.reset_head(num_classes=1)
return net
def __init__(self, num_classes=11, checkpoint=None):
super(EfficientDet6, self).__init__()
config = get_efficientdet_config('tf_efficientdet_d6')
config.image_size = [512, 512]
config.norm_kwargs = dict(eps=.001, momentum=.01)
config.soft_nms = True
config.label_smoothing = 0.1
config.mean = [0.46009655, 0.43957878, 0.41827092]
config.std = [0.2108204, 0.20766491, 0.21656131]
net = EfficientDet(config, pretrained_backbone=False)
if checkpoint == None:
checkpoint = torch.load(
'./effdet_model/tf_efficientdet_d6_52-4eda3773.pth')
net.load_state_dict(checkpoint)
net.reset_head(num_classes=num_classes)
net.class_net = HeadNet(config, num_outputs=config.num_classes)
self.model = DetBenchTrain(net, config)
def Train_Model(which_model='effi5',
img_size=1280,
weight_dir): # 모델 및 pretrained weight 불러오기
if which_model == 'effi5':
config = get_efficientdet_config(
'tf_efficientdet_d5') # tf_effi5 model structure
config.image_size = (img_size, img_size)
config.norm_kwargs = dict(eps=.001, momentum=.01)
# net = create_model_from_config(config, bench_task='train',
# num_classes=32, checkpoint_path='../efficientdet-pytorch/tf_efficientdet_d5-ef44aea8.pth')
net = EfficientDet(config, pretrained_backbone=False)
pretrained_weight = torch.load(weight_dir +
'/tf_efficientdet_d5-ef44aea8.pth')
net.load_state_dict(pretrained_weight)
# 기본 net에다가 pretrained weight을 불러온 후 class를 바꿔야함 (reset)
elif which_model == 'effi4':
config = get_efficientdet_config(
'tf_efficientdet_d4') # tf_effi4 model structure
config.image_size = (img_size, img_size)
config.norm_kwargs = dict(eps=.001, momentum=.01)
net = EfficientDet(config, pretrained_backbone=False)
pretrained_weight = torch.load(weight_dir +
'/tf_efficientdet_d4-5b370b7a.pth')
net.load_state_dict(pretrained_weight)
elif which_model == 'effi6':
config = get_efficientdet_config(
'tf_efficientdet_d6') # tf_effi4 model structure
config.image_size = (img_size, img_size)
config.norm_kwargs = dict(eps=.001, momentum=.01)
net = EfficientDet(config, pretrained_backbone=False)
pretrained_weight = torch.load(weight_dir +
'/tf_efficientdet_d6-51cb0132.pth')
net.load_state_dict(pretrained_weight)
net.reset_head(num_classes=32)
net.class_net = HeadNet(config, num_outputs=config.num_classes)
return DetBenchTrain(net, config)
