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 get_efficientdet(checkpoint_path):
config = get_efficientdet_config('efficientdet_d1')
net = EfficientDet(config, pretrained_backbone=False)
checkpoint = torch.load(checkpoint_path)
net.load_state_dict(checkpoint, strict=False)
net = DetBenchPredict(net)
return net
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 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 __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 get_net():
config = get_efficientdet_config('efficientdet_d0')
config.num_classes = 29
config.image_size = [512, 512]
net = EfficientDet(config, pretrained_backbone=False)
checkpoint = torch.load(
'effdet_checkpoints/2021-02-28 19_55_41/last-checkpoint.pth')
net.load_state_dict(checkpoint['model_state_dict'])
return DetBenchPredict(net)
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 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 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('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 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_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 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_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 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_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 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_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 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 __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():
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 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_effdet_train_hotstart(checkpoint_path: str):
config = get_efficientdet_config("tf_efficientdet_d5")
model = EfficientDet(config, pretrained_backbone=False)
config.num_classes = 1
config.image_size = 512
model.class_net = HeadNet(
config,
num_outputs=config.num_classes,
norm_kwargs=dict(eps=0.001, momentum=0.01),
)
checkpoint = torch.load(checkpoint_path)
model.load_state_dict(checkpoint["model_state_dict"])
del checkpoint
gc.collect()
model = DetBenchTrain(model, config).train()
return model
def get_model_(variant, model_dir):
config = get_efficientdet_config(f"tf_efficientdet_{variant}")
net = EfficientDet(config, pretrained_backbone=False)
# checkpoint = torch.load('/content/wheat_efficientdet/pretrained_models/efficientdet_d0-d92fd44f.pth')
if variant == "d0":
checkpoint_path = f"{model_dir}/efficientdet_d0-d92fd44f.pth"
checkpoint = torch.load(checkpoint_path)
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=0.001, momentum=0.01),
)
return DetBenchTrain(net, 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()
def get_model(variant, model_dir, load_path):
config = get_efficientdet_config(f"tf_efficientdet_{variant}")
net = EfficientDet(config, pretrained_backbone=False)
if not load_path:
print("here")
if variant == "d0":
checkpoint_path = f"{model_dir}/efficientdet_d0-d92fd44f.pth"
elif variant == "d1":
checkpoint_path = f"{model_dir}/efficientdet_d1-4c7ebaf2.pth"
elif variant == "d2":
checkpoint_path = f"{model_dir}/efficientdet_d2-cb4ce77d.pth"
elif variant == "d3":
checkpoint_path = f"{model_dir}/efficientdet_d3-b0ea2cbc.pth"
elif variant == "d4":
checkpoint_path = f"{model_dir}/efficientdet_d4-5b370b7a.pth"
elif variant == "d5":
checkpoint_path = f"{model_dir}/efficientdet_d5-ef44aea8.pth"
elif variant == "d6":
checkpoint_path = f"{model_dir}/efficientdet_d6-51cb0132.pth"
elif variant == "d6":
checkpoint_path = f"{model_dir}/efficientdet_d7-f05bf714.pth"
checkpoint = torch.load(checkpoint_path)
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=0.001, momentum=0.01),
)
if load_path:
print("here")
checkpoint = torch.load(checkpoint_path)
net.load_state_dict(checkpoint["model_state_dict"])
return DetBenchTrain(net, config)
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 get_train_model(config_name,
img_size,
model_ckpt=None,
useGN=False,
light=True):
config = get_efficientdet_config(config_name)
model = EfficientDet(config, pretrained_backbone=True)
config.num_classes = 1
config.image_size = img_size
model.class_net = HeadNet(
config,
num_outputs=config.num_classes,
norm_kwargs=dict(eps=0.001, momentum=0.01),
)
if useGN is True:
model = convert_layers(model,
nn.BatchNorm2d,
nn.GroupNorm,
True,
num_groups=2)
# Replace BatchNorm with GroupNorm
if model_ckpt is not None:
if light is True:
count = 0
state_dict = torch.load(model_ckpt)["state_dict"]
new_state_dict = OrderedDict()
for key, value in state_dict.items():
if key.startswith("model.model."):
new_key = reduce(lambda a, b: a + "." + b,
key.split(".")[2:])
if new_key in model.state_dict():
new_state_dict[new_key] = value
count += 1
model.load_state_dict(new_state_dict)
print(f"loaded {count} keys")
else:
model.load_state_dict(torch.load(model_ckpt)["state_dict"])
return DetBenchTrain(model, config)
