def __init__(self, classes=1000, load_backbone_weights=False, config_file=None, weights_file=None): """ load the model and the sequential head so that the backbone can be applied for classification Model tested on ImageNet Tiny Args: classes: integer for how many classes can be predicted load_backbone_weights: bool, if true we will auto load the original darnet weights to use in the model config_file: str path for the location of the configuration file to use when decoding darknet weights weights_file: str path with the file containing the dark net weights """ super(DarkNet53, self).__init__() self.backbone = Backbone_Builder("darknet53") self.head = ks.Sequential([ ks.layers.GlobalAveragePooling2D(), ks.layers.Dense(classes, activation="sigmoid") ]) if load_backbone_weights: if config_file is None: config_file = download('yolov3.cfg') if weights_file is None: weights_file = download('yolov3.weights') full_model = DarkNetConverter.read(config_file, weights_file) encoder, decoder = split_converter(full_model, 76) load_weights_dnBackbone(self.backbone, encoder) return
def load_weights_from_dn(self, dn2tf_backbone=True, dn2tf_head=True, config_file=None, weights_file=None): """ load the entire Yolov3 Model for tensorflow example: load yolo with darknet wieghts for backbone model = Yolov3() model.build(input_shape = (1, 416, 416, 3)) model.load_weights_from_dn(dn2tf_backbone = True, dn2tf_head = True) to be implemented example: load custom back bone weigths example: load custom head weigths example: load back bone weigths from tensorflow (our training) example: load head weigths from tensorflow (our training) Args: dn2tf_backbone: bool, if true it will load backbone weights for yolo v3 from darknet .weights file dn2tf_head: bool, if true it will load head weights for yolo v3 from darknet .weights file config_file: str path for the location of the configuration file to use when decoding darknet weights weights_file: str path with the file containing the dark net weights """ if not self._built: self.build(self._input_shape) if dn2tf_backbone or dn2tf_neck or dn2tf_head: if config_file is None: config_file = download(self._model_name + '.cfg') if weights_file is None: weights_file = download(self._model_name + '.weights') list_encdec = DarkNetConverter.read(config_file, weights_file) encoder, neck, decoder = split_converter( list_encdec, self._encoder_decoder_split_location, 138) if dn2tf_backbone: #load_weights_dnBackbone(self._backbone, encoder, mtype = self._backbone_name) load_weights_backbone(self._backbone, encoder) self._backbone.trainable = False if dn2tf_head: load_weights_backbone(self._neck, neck) self._neck.trainable = False load_weights_v4head(self._head, decoder) self._head.trainable = False return
def initialize(self, model: tf.keras.Model): if self.task_config.load_darknet_weights: from yolo.utils import DarkNetConverter from yolo.utils._darknet2tf.load_weights import split_converter from yolo.utils._darknet2tf.load_weights2 import load_weights_backbone from yolo.utils._darknet2tf.load_weights2 import load_head from yolo.utils._darknet2tf.load_weights2 import load_weights_prediction_layers from yolo.utils.downloads.file_manager import download weights_file = self.task_config.model.darknet_weights_file config_file = self.task_config.model.darknet_weights_cfg if ('cache' not in weights_file and 'cache' not in config_file): list_encdec = DarkNetConverter.read(config_file, weights_file) else: import os path = os.path.abspath('cache') if (not os.path.isdir(path)): os.mkdir(path) cfg = f"{path}/cfg/{config_file.split('/')[-1]}" if not os.path.isfile(cfg): download(config_file.split('/')[-1]) wgt = f"{path}/weights/{weights_file.split('/')[-1]}" if not os.path.isfile(wgt): download(weights_file.split('/')[-1]) list_encdec = DarkNetConverter.read(cfg, wgt) splits = model.backbone._splits if 'neck_split' in splits.keys(): encoder, neck, decoder = split_converter( list_encdec, splits['backbone_split'], splits['neck_split']) else: encoder, decoder = split_converter(list_encdec, splits['backbone_split']) neck = None load_weights_backbone(model.backbone, encoder) model.backbone.trainable = False if self.task_config.darknet_load_decoder: if neck != None: load_weights_backbone(model.decoder.neck, neck) model.decoder.neck.trainable = False cfgheads = load_head(model.decoder.head, decoder) model.decoder.head.trainable = False load_weights_prediction_layers(cfgheads, model.head) model.head.trainable = False else: """Loading pretrained checkpoint.""" if not self.task_config.init_checkpoint: return ckpt_dir_or_file = self.task_config.init_checkpoint if tf.io.gfile.isdir(ckpt_dir_or_file): ckpt_dir_or_file = tf.train.latest_checkpoint(ckpt_dir_or_file) # Restoring checkpoint. if self.task_config.init_checkpoint_modules == 'all': ckpt = tf.train.Checkpoint(**model.checkpoint_items) status = ckpt.restore(ckpt_dir_or_file) status.assert_consumed() elif self.task_config.init_checkpoint_modules == 'backbone': ckpt = tf.train.Checkpoint(backbone=model.backbone) status = ckpt.restore(ckpt_dir_or_file) status.expect_partial().assert_existing_objects_matched() else: assert "Only 'all' or 'backbone' can be used to initialize the model." logging.info('Finished loading pretrained checkpoint from %s', ckpt_dir_or_file)