def __init__(self, cfg_path,
weight_path,
resume,
gpu_id
):
init_seeds(0)
self.device = gpu.select_device(gpu_id)
self.start_epoch = 0
self.best_mAP = 0.
self.epochs = TRAIN["EPOCHS"]
self.weight_path = weight_path
self.multi_scale_train = TRAIN["MULTI_SCALE_TRAIN"]
self.train_dataset = data.VocDataset(anno_file_type="train", img_size=TRAIN["TRAIN_IMG_SIZE"])
self.train_dataloader = DataLoader(self.train_dataset, batch_size=TRAIN["BATCH_SIZE"],
num_workers=TRAIN["NUMBER_WORKERS"], shuffle=True)
self.yolov3 = Darknet(cfg_path=cfg_path, img_size=TRAIN["TRAIN_IMG_SIZE"]).to(self.device)
self.yolov3.apply(tools.weights_init_normal)
self.optimizer = optim.SGD(self.yolov3.parameters(), lr=TRAIN["LR_INIT"], momentum=TRAIN["MOMENTUM"],
weight_decay=TRAIN["WEIGHT_DECAY"])
# self.optimizer = optim.Adam(self.yolov3.parameters(), lr = lr_init, weight_decay=0.9995)
self.criterion = YoloV3Loss(anchors=MODEL["ANCHORS"], strides=MODEL["STRIDES"],
iou_threshold_loss=TRAIN["IOU_THRESHOLD_LOSS"])
self.__load_model_weights(weight_path, resume)
self.scheduler = lr_scheduler.MultiStepLR(self.optimizer, milestones=[25, 40], gamma=0.1,
last_epoch=self.start_epoch - 1)
def __init__(self, model, weight_path, resume, gpu_id):
init_seeds(0)
self.model = model
self.device = gpu.select_device(gpu_id)
self.start_epoch = 0
self.best_mAP = 0.
self.epochs = cfg.TRAIN["EPOCHS"]
self.weight_path = weight_path
self.multi_scale_train = cfg.TRAIN["MULTI_SCALE_TRAIN"]
self.train_dataset = data.VocDataset(
anno_file_type="train", img_size=cfg.TRAIN["TRAIN_IMG_SIZE"])
self.train_dataloader = DataLoader(
self.train_dataset,
batch_size=cfg.TRAIN["BATCH_SIZE"],
num_workers=cfg.TRAIN["NUMBER_WORKERS"],
shuffle=True)
if self.model == 's':
self.yolov3 = Yolov3_S().to(self.device)
self.model_postfix = '_s'
elif self.model == 'l':
self.yolov3 = Yolov3_L().to(self.device)
self.model_postfix = '_l'
else:
self.yolov3 = Yolov3().to(self.device)
self.model_postfix = '_m'
# self.yolov3.apply(tools.weights_init_normal)
self.optimizer = optim.SGD(self.yolov3.parameters(),
lr=cfg.TRAIN["LR_INIT"],
momentum=cfg.TRAIN["MOMENTUM"],
weight_decay=cfg.TRAIN["WEIGHT_DECAY"])
#self.optimizer = optim.Adam(self.yolov3.parameters(), lr=cfg.TRAIN["LR_INIT"], weight_decay=0.9995)
self.criterion = YoloV3Loss(
anchors=cfg.MODEL["ANCHORS"],
strides=cfg.MODEL["STRIDES"],
iou_threshold_loss=cfg.TRAIN["IOU_THRESHOLD_LOSS"])
self.__load_model_weights(weight_path, resume)
self.scheduler = cosine_lr_scheduler.CosineDecayLR(
self.optimizer,
T_max=self.epochs * len(self.train_dataloader),
lr_init=cfg.TRAIN["LR_INIT"],
lr_min=cfg.TRAIN["LR_END"],
warmup=cfg.TRAIN["WARMUP_EPOCHS"] * len(self.train_dataloader))