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_net(checkpoint_path):
config = get_efficientdet_config('tf_efficientdet_d7')
net = EfficientDet(config, pretrained_backbone=False)
config.num_classes = 1
config.image_size=512
net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))
print(f"Load model {checkpoint_path}")
checkpoint = torch.load(checkpoint_path)
net.cuda()
new_state_dict = OrderedDict()
for k, v in checkpoint['model_state_dict'].items():
if "anchors" in k:
print("Ignore: ",k)
continue
name = re.sub("model.",'',k) if k.startswith('model') else k
new_state_dict[name] = v
# net.load_state_dict(checkpoint['model_state_dict'])
net.load_state_dict(new_state_dict)
del checkpoint
gc.collect()
net = DetBenchPredict(net, config)
net.eval()
return net.cuda()
def Load_Model(checkpoint_path, which_model='effi5', img_size=1280):
if which_model == 'effi5':
config = get_efficientdet_config(
'tf_efficientdet_d5') # tf_effi5 model structure
elif which_model == 'effi4':
config = get_efficientdet_config(
'tf_efficientdet_d4') # tf_effi4 model structure
elif which_model == 'effi6':
config = get_efficientdet_config(
'tf_efficientdet_d6') # tf_effi4 model structure
config.image_size = (img_size, img_size)
config.num_classes = 32
config.norm_kwargs = dict(eps=.001, momentum=.01)
net = EfficientDet(config, pretrained_backbone=False)
net.class_net = HeadNet(config, num_outputs=config.num_classes)
ckp = torch.load(checkpoint_path)
net.load_state_dict(ckp['model_state_dict'])
del ckp
net = DetBenchPredict(net)
net.eval()
return net
def get_model_detection_efficientdet(model_name,
num_classes,
target_dim,
freeze_batch_norm=False):
print("Using EffDet detection model")
config = effdet.get_efficientdet_config(model_name)
efficientDetModel = EfficientDet(config, pretrained_backbone=False)
load_pretrained(efficientDetModel, config.url)
import omegaconf
with omegaconf.read_write(config):
config.num_classes = num_classes
# config.image_size = target_dim
efficientDetModel.class_net = HeadNet(config, num_outputs=num_classes)
if freeze_batch_norm:
# we only freeze BN layers in backbone and the BiFPN
print("Freezing batch normalization weights")
freeze_bn(efficientDetModel.backbone)
with omegaconf.read_write(efficientDetModel.config):
efficientDetModel.config.num_classes = num_classes
# print(DetBenchTrain(efficientDetModel, config))
return DetBenchTrain(efficientDetModel, config)
def init_settings(self):
""" initialization of data and model variables """
# load nn model
config = get_efficientdet_config('tf_efficientdet_d5')
net = EfficientDet(config, pretrained_backbone=False)
checkpoint = torch.load(
'input/efficientdet/efficientdet_d5-ef44aea8.pth')
net.load_state_dict(checkpoint)
config.num_classes = 1
config.image_size = 512
net.class_net = HeadNet(config,
num_outputs=config.num_classes,
norm_kwargs=dict(eps=.001, momentum=.01))
self.model = DetBenchMulty(net, config)
if args.weights: # load prev weight and continue train process
checkpoint = torch.load(os.path.join(wdir, args.weights))
try:
self.model.load_state_dict(checkpoint['model_state_dict'])
except RuntimeError:
state_dict = torch.load(os.path.join(wdir, args.weights))
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict['model_state_dict'].items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
# load params
self.model.load_state_dict(new_state_dict)
def get_net():
config = get_efficientdet_config('tf_efficientdet_d0')
net = EfficientDet(config, pretrained_backbone=True)
config.num_classes = 14
config.image_size = 512
net.class_net = HeadNet(config,
num_outputs=config.num_classes,
norm_kwargs=dict(eps=.001, momentum=.01))
return DetBenchTrain(net, config)
def get_net(model_name):
config = get_efficientdet_config(model_name)
config.norm_kwargs = dict(eps=.001, momentum=.01)
config.num_classes = 29
# config.image_size = [640, 640]
net = EfficientDet(config, pretrained_backbone=True)
net.class_net = HeadNet(config, num_outputs=config.num_classes)
return DetBenchTrain(net, config)
def __init__(self, model, num_class):
super(EfficientDetCus, self).__init__()
config = get_efficientdet_config(f'tf_efficientdet_{model}')
config.num_classes = num_class
config.image_size = [TRAIN_SIZE, TRAIN_SIZE]
model = EfficientDet(config=config, pretrained_backbone=False)
model.class_net = HeadNet(config, num_outputs=config.num_classes)
self.model = DetBenchTrain(model, config)
def get_net():
config = get_efficientdet_config('tf_efficientdet_d5')
net = EfficientDet(config, pretrained_backbone=False)
checkpoint = torch.load(r'/content/gdrive/My Drive/Colab Notebooks/globalwheat/input/efficientdet-pytorch-master/efficientdet_d5-ef44aea8.pth')
net.load_state_dict(checkpoint)
config.num_classes = 1
config.image_size = 512
net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))
return DetBenchTrain(net, config)
def get_net():
config = get_efficientdet_config('tf_efficientdet_d5')
net = EfficientDet(config, pretrained_backbone=False)
checkpoint = torch.load('./input/efficientdet/efficientdet_d5-ef44aea8.pth')
net.load_state_dict(checkpoint)
config.num_classes = 1
config.image_size = 512
net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))
return DetBenchTrain(net, config)
def model(model_name: str,
num_classes: int,
img_size: int,
pretrained: bool = True) -> nn.Module:
"""Creates the efficientdet model specified by `model_name`.
The model implementation is by Ross Wightman, original repo
[here](https://github.com/rwightman/efficientdet-pytorch).
# Arguments
model_name: Specifies the model to create. For pretrained models, check
[this](https://github.com/rwightman/efficientdet-pytorch#models) table.
num_classes: Number of classes of your dataset (including background).
img_size: Image size that will be fed to the model. Must be squared and
divisible by 64.
pretrained: If True, use a pretrained backbone (on COCO).
# Returns
A PyTorch model.
"""
config = get_efficientdet_config(model_name=model_name)
net = EfficientDet(config, pretrained_backbone=False)
if pretrained:
if not config.url:
raise RuntimeError(f"No pretrained weights for {model_name}")
state_dict = torch.hub.load_state_dict_from_url(
config.url, map_location=torch.device("cpu"))
net.load_state_dict(state_dict)
config.num_classes = num_classes
config.image_size = img_size
net.class_net = HeadNet(
config,
num_outputs=num_classes,
norm_kwargs=dict(eps=0.001, momentum=0.01),
)
# TODO: Break down param groups for backbone
def param_groups_fn(model: nn.Module) -> List[List[nn.Parameter]]:
unwrapped = unwrap_bench(model)
layers = [
unwrapped.backbone,
unwrapped.fpn,
nn.Sequential(unwrapped.class_net, unwrapped.box_net),
]
param_groups = [list(layer.parameters()) for layer in layers]
check_all_model_params_in_groups2(model, param_groups)
return param_groups
model_bench = DetBenchTrain(net, config)
model_bench.param_groups = MethodType(param_groups_fn, model_bench)
return model_bench
def get_net(ckpt_path = None):
config = get_efficientdet_config('tf_efficientdet_d5')
net = EfficientDet(config, pretrained_backbone=False)
if(ckpt_path !=None):
checkpoint = torch.load('../input/d5-pseudo-aug-best/best-checkpoint-026epoch-2.bin')
config.num_classes = 1
config.image_size = 512
net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))
if(ckpt_path !=None):
net.load_state_dict(checkpoint['model_state_dict'])
return DetBenchTrain(net, config)
def get_net():
config = get_efficientdet_config(Preprocess_config().model)
net = EfficientDet(config, pretrained_backbone=True)
#checkpoint = torch.load('../input/efficientdet/efficientdet_d5-ef44aea8.pth')
#net.load_state_dict(checkpoint)
config.num_classes = 1
config.image_size = 1024
net.class_net = HeadNet(config,
num_outputs=config.num_classes,
norm_kwargs=dict(eps=.001, momentum=.01))
return DetBenchTrain(net, config)
def get_net():
config = get_efficientdet_config('tf_efficientdet_d5')
net = EfficientDet(config, pretrained_backbone=False)
checkpoint = torch.load(args["weights_path"])
net.load_state_dict(checkpoint)
config.num_classes = 1
config.image_size = args["image_size"]
net.class_net = HeadNet(config,
num_outputs=config.num_classes,
norm_kwargs=dict(eps=.001, momentum=.01))
return DetBenchTrain(net, config)
def get_net():
config = get_efficientdet_config('tf_efficientdet_d6')
net = EfficientDet(config, pretrained_backbone=False)
config.num_classes = 1
config.image_size = 1024
net.class_net = HeadNet(config,
num_outputs=config.num_classes,
norm_kwargs=dict(eps=.001, momentum=.01))
checkpoint = torch.load('pretrained/efficientdet_d6-51cb0132.pth')
load_state_dict(net, checkpoint)
return DetBenchTrain(net, config)
def set_train_effdet(config,
num_classes: int = 1,
device: torch.device = 'cuda:0'):
"""Init EfficientDet to train mode"""
model = EfficientDet(config, pretrained_backbone=False)
model.class_net = HeadNet(config,
num_outputs=num_classes,
norm_kwargs=dict(eps=.001, momentum=.01))
model = DetBenchTrain(model, config)
model = model.train()
return model.to_device(device)
def get_effdet_train(model_path: str):
config = get_efficientdet_config("tf_efficientdet_d5")
net = EfficientDet(config, pretrained_backbone=False)
checkpoint = torch.load(model_path)
config.num_classes = 1
config.image_size = 512
net.load_state_dict(checkpoint)
net.class_net = HeadNet(
config,
num_outputs=config.num_classes,
norm_kwargs=dict(eps=0.001, momentum=0.01),
)
return DetBenchTrain(net, config).train()
def load_net(checkpoint_path):
config = get_efficientdet_config('tf_efficientdet_d5')
net = EfficientDet(config, pretrained_backbone=False)
config.num_classes = 2
config.image_size = 512
net.class_net = HeadNet(config,
num_outputs=config.num_classes,
norm_kwargs=dict(eps=.001, momentum=.01))
checkpoint = torch.load(checkpoint_path)
net.load_state_dict(checkpoint['model_state_dict'])
net = DetBenchEval(net, config)
net.eval()
return net.cuda()
def get_net():
config = get_efficientdet_config('tf_efficientdet_d6')
net = EfficientDet(config, pretrained_backbone=False)
config.num_classes = 1
config.image_size = 1024
net.class_net = HeadNet(config,
num_outputs=config.num_classes,
norm_kwargs=dict(eps=.001, momentum=.01))
checkpoint = torch.load(
'effdet6-baseline-1024-4x8-sa-fold0/best-checkpoint-052epoch.bin')
load_state_dict(net, checkpoint['model_state_dict'])
return DetBenchTrain(net, config)
def get_net():
config = get_efficientdet_config('tf_efficientdet_d5')
net = EfficientDet(config, pretrained_backbone=False)
checkpoint = torch.load(
'/home/eragon/Documents/scripts/efficientdet-pytorch/checkpoints/tf_efficientdet_d5_51-c79f9be6.pth'
)
net.load_state_dict(checkpoint)
config.num_classes = 1
config.image_size = 512
net.class_net = HeadNet(config,
num_outputs=config.num_classes,
norm_kwargs=dict(eps=.001, momentum=.01))
return DetBenchTrain(net, config)
def model(model_name: str,
num_classes: int,
img_size: int,
pretrained: bool = True) -> nn.Module:
""" Creates the model specific by model_name
Args:
model_name (str): Specifies the model to create, available options are: TODO
num_classes (int): Number of classes of your dataset (including background)
pretrained (int): If True, use a pretrained backbone (on COCO)
Returns:
nn.Module: The requested model
"""
config = get_efficientdet_config(model_name=model_name)
net = EfficientDet(config, pretrained_backbone=False)
if pretrained:
if not config.url:
raise RuntimeError(f"No pretrained weights for {model_name}")
state_dict = torch.hub.load_state_dict_from_url(
config.url, map_location=torch.device("cpu"))
net.load_state_dict(state_dict)
config.num_classes = num_classes
config.image_size = img_size
net.class_net = HeadNet(
config,
num_outputs=num_classes,
norm_kwargs=dict(eps=0.001, momentum=0.01),
)
# TODO: Break down param groups for backbone
def param_groups_fn(model: nn.Module) -> List[List[nn.Parameter]]:
unwrapped = unwrap_bench(model)
layers = [
unwrapped.backbone,
unwrapped.fpn,
nn.Sequential(unwrapped.class_net, unwrapped.box_net),
]
param_groups = [list(layer.parameters()) for layer in layers]
check_all_model_params_in_groups2(model, param_groups)
return param_groups
model_bench = DetBenchTrain(net, config)
model_bench.param_groups = MethodType(param_groups_fn, model_bench)
return model_bench
def get_net():
config = get_efficientdet_config('tf_efficientdet_d5')
net = EfficientDet(config, pretrained_backbone=False)
# config.num_classes = 1
# config.image_size = 512
# net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))
# checkpoint = torch.load(r'D:\Workspace\efficientdet-pytorch-master\effdet5-cutmix-augmix\last-checkpoint.bin')
# net.load_state_dict(checkpoint['model_state_dict'])
checkpoint = torch.load(r'C:\Users\Aministration\Downloads\efficientdet_d5-ef44aea8.pth')
net.load_state_dict(checkpoint)
config.num_classes = 1
config.image_size = 512
net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))
return DetBenchTrain(net, config)
def get_net():
config = get_efficientdet_config('tf_efficientdet_d7')
net = EfficientDet(config, pretrained_backbone=False)
# config.num_classes = 1
# config.image_size = 512
# net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))
# checkpoint = torch.load(r'./effdet5-cutmix-augmix1/last-checkpoint-folder1.bin')
# net.load_state_dict(checkpoint['model_state_dict'])
checkpoint = torch.load(r'./tf_efficientdet_d7_53-6d1d7a95.pth')
net.load_state_dict(checkpoint)
config.num_classes = 1
config.image_size = 512
net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))
return DetBenchTrain(net, config)
def load_net():
config = get_efficientdet_config('tf_efficientdet_d0')
net = EfficientDet(config, pretrained_backbone=False)
config.num_classes = 1
config.image_size=512
net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))
checkpoint = torch.load(DIR_PATH + '/models/effdet_trained.pth')
net.load_state_dict(checkpoint["model_state_dict"])
del checkpoint
gc.collect()
net = DetBenchEval(net, config)
net.eval();
return net.cuda()
def get_net():
# 模型的配置，这个返回的是一个字典
config = get_efficientdet_config('tf_efficientdet_d5')
config.num_classes = 1
config.image_size = 512
# 根据上面的配置生成网络
net = EfficientDet(config, pretrained_backbone=False)
# 加载gwd预训练模型
checkpoint = torch.load(r'../input/eff_checkpoint/bestv2.bin')
net.load_state_dict(checkpoint['model_state_dict'])
# 加載coco預訓練模型
# norm_kwargs 设置的是 BATCHNORM2D 的参数
net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))
return DetBenchTrain(net, config)
def __init__(self, model_weight, num_class):
super(EfficientDetPred, self).__init__()
config = get_efficientdet_config(f'tf_efficientdet_{MODEL_USE}')
config.num_classes = num_class
config.image_size = [TRAIN_SIZE, TRAIN_SIZE]
model = EfficientDet(config, pretrained_backbone=False)
model.class_net = HeadNet(config, num_outputs=config.num_classes)
new_keys = model.state_dict().keys()
values = torch.load(model_weight, map_location=lambda storage, loc: storage).values()
model.load_state_dict(OrderedDict(zip(new_keys, values)))
self.model = DetBenchPredict(model)
del new_keys, values
gc.collect()
def load_net(checkpoint_path):
config = get_efficientdet_config('tf_efficientdet_d5')
net = EfficientDet(config, pretrained_backbone=False)
config.num_classes = 1
config.image_size = 512
net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))
checkpoint = torch.load(checkpoint_path, map_location=DEVICE)
net.load_state_dict(checkpoint['model_state_dict'])
del checkpoint
gc.collect()
net = DetBenchTrain(net, config)
return net.cuda()
def set_eval_effdet(checkpoint_path: str,
config,
num_classes: int = 1,
device: torch.device = 'cuda:0'):
"""Init EfficientDet to validation mode"""
net = EfficientDet(config, pretrained_backbone=False)
net.class_net = HeadNet(config,
num_outputs=num_classes,
norm_kwargs=dict(eps=.001, momentum=.01))
checkpoint = torch.load(checkpoint_path)
net.load_state_dict(checkpoint)
net = DetBenchEval(net, config)
net = net.eval()
return net.to_device(device)
def get_train_model(config_name='tf_efficientdet_d0', model_ckpt=None):
config = get_efficientdet_config(config_name)
model = EfficientDet(config, pretrained_backbone=True)
#layers_to_train = ['blocks.6','blocks.5', 'blocks.4', 'blocks.3', 'blocks.2', 'blocks.1', 'conv_stem'][:5]
#for name, parameter in model.backbone.named_parameters():
# if all([not name.startswith(layer) for layer in layers_to_train]):
# parameter.requires_grad_(False)
# print(f'layer {name} frozen')
config.num_classes = 1
config.image_size = IMG_SIZE
model.class_net = HeadNet(config,
num_outputs=config.num_classes,
norm_kwargs=dict(eps=.001, momentum=.01))
if model_ckpt is not None:
model.load_state_dict(torch.load(model_ckpt)['model_state_dict'])
return DetBenchTrain(model, config)
def get_test_net(best_weigth):
config = get_efficientdet_config('tf_efficientdet_d5')
net = EfficientDet(config, pretrained_backbone=False)
config.num_classes = 1
config.image_size = 512
net.class_net = HeadNet(config,
num_outputs=config.num_classes,
norm_kwargs=dict(eps=.001, momentum=.01))
# checkpoint = torch.load(checkpoint_path)
net.load_state_dict(best_weigth)
# del checkpoint
gc.collect()
net = DetBenchEval(net, config)
net = net.train()
return net.cuda()
