def __init__(self, weight_path, resume, gpu_id, accumulate, fp_16):
init_seeds(0)
self.fp_16 = fp_16
self.device = gpu.select_device(gpu_id)
self.start_epoch = 0
self.best_mAP = 0.
self.accumulate = accumulate
self.epochs = cfg.TRAIN["EPOCHS"]
self.weight_path = weight_path
self.multi_scale_train = cfg.TRAIN["MULTI_SCALE_TRAIN"]
self.train_dataset = data.Build_Dataset(anno_file_type="train", img_size=cfg.TRAIN["TRAIN_IMG_SIZE"])
print('train img size is {}'.format(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, pin_memory=True
)
self.yolov4 = Build_Model().to(self.device)
self.optimizer = optim.SGD(self.yolov4.parameters(), lr=cfg.TRAIN["LR_INIT"],
momentum=cfg.TRAIN["MOMENTUM"], weight_decay=cfg.TRAIN["WEIGHT_DECAY"])
self.criterion = YoloV4Loss(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))
def __init__(self,
weight_path=None,
resume=False,
gpu_id=0,
accumulate=1,
fp_16=False):
init_seeds(0)
self.fp_16 = fp_16
self.device = gpu.select_device(gpu_id)
self.start_epoch = 0
self.best_mAP = 0.0
self.accumulate = accumulate
self.weight_path = weight_path
self.multi_scale_train = cfg.TRAIN["MULTI_SCALE_TRAIN"]
self.showatt = cfg.TRAIN["showatt"]
if self.multi_scale_train:
print("Using multi scales training")
else:
print("train img size is {}".format(cfg.TRAIN["TRAIN_IMG_SIZE"]))
self.train_dataset = data.Build_Dataset(
anno_file_type="train", img_size=cfg.TRAIN["TRAIN_IMG_SIZE"])
self.epochs = (cfg.TRAIN["YOLO_EPOCHS"] if cfg.MODEL_TYPE["TYPE"]
== "YOLOv4" else cfg.TRAIN["Mobilenet_YOLO_EPOCHS"])
self.eval_epoch = (30 if cfg.MODEL_TYPE["TYPE"] == "YOLOv4" else 50)
self.train_dataloader = DataLoader(
self.train_dataset,
batch_size=cfg.TRAIN["BATCH_SIZE"],
num_workers=cfg.TRAIN["NUMBER_WORKERS"],
shuffle=True,
pin_memory=True,
)
self.yolov4 = Build_Model(weight_path=weight_path,
resume=resume,
showatt=self.showatt).to(self.device)
self.optimizer = optim.SGD(
self.yolov4.parameters(),
lr=cfg.TRAIN["LR_INIT"],
momentum=cfg.TRAIN["MOMENTUM"],
weight_decay=cfg.TRAIN["WEIGHT_DECAY"],
)
self.criterion = YoloV4Loss(
anchors=cfg.MODEL["ANCHORS"],
strides=cfg.MODEL["STRIDES"],
iou_threshold_loss=cfg.TRAIN["IOU_THRESHOLD_LOSS"],
)
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),
)
if resume:
self.__load_resume_weights(weight_path)
help='batches to accumulate before optimizing')
parser.add_argument('--fp_16',
type=bool,
default=False,
help='whither to use fp16 precision')
parser.add_argument('--exp_name',
type=str,
default='debug',
help='log experiment name')
parser.add_argument('--ckpt',
type=str,
default=None,
help='model checkpoint') #weight/darknet53_448.weights
opt = parser.parse_args()
train_dataset = data.Build_Dataset(anno_file_type="train",
img_size=cfg.TRAIN["TRAIN_IMG_SIZE"])
test_dataset = data.Build_Dataset(anno_file_type="test",
img_size=cfg.VAL["TEST_IMG_SIZE"])
train_dataloader = DataLoader(
train_dataset,
batch_size=1, #cfg.TRAIN["BATCH_SIZE"],
num_workers=cfg.TRAIN["NUMBER_WORKERS"],
shuffle=True,
pin_memory=True)
test_dataloader = DataLoader(
test_dataset,
batch_size=1, #cfg.VAL["BATCH_SIZE"],
num_workers=cfg.VAL["NUMBER_WORKERS"],
shuffle=False,
pin_memory=True)
def __init__(self,
weight_path=None,
resume: bool = False,
gpu_id: int = 0,
accumulate: bool = True,
fp_16: bool = False):
# PYTHON HASH SEED
init_seeds(0)
# device
self.fp_16: bool = fp_16
self.device: torch.device = gpu.select_device(gpu_id)
self.start_epoch: int = 0
self.best_mAP: float = 0.0 # not sure why this is necessary...
self.accumulate: bool = accumulate
self.weight_path: Path = weight_path
self.multi_scale_train: bool = cfg.TRAIN["MULTI_SCALE_TRAIN"]
# Show attention modification?
self.showatt = cfg.TRAIN["showatt"]
# Multi-scale training status
if self.multi_scale_train:
print("Using multi scales training")
else:
print(f"train img size is {cfg.TRAIN['TRAIN_IMG_SIZE']}")
# Build Dataset using helper function.
self.train_dataset = data.Build_Dataset(
anno_file_type="train", img_size=cfg.TRAIN["TRAIN_IMG_SIZE"])
self.epochs = (cfg.TRAIN["YOLO_EPOCHS"] if cfg.MODEL_TYPE["TYPE"]
== "YOLOv4" else cfg.TRAIN["Mobilenet_YOLO_EPOCHS"])
self.eval_epoch = (30 if cfg.MODEL_TYPE["TYPE"] == "YOLOv4" else 50)
self.train_dataloader = DataLoader(
self.train_dataset,
batch_size=cfg.TRAIN["BATCH_SIZE"],
num_workers=cfg.TRAIN["NUMBER_WORKERS"],
shuffle=True,
pin_memory=True,
)
self.yolov4 = Build_Model(weight_path=weight_path,
resume=resume,
showatt=self.showatt).to(self.device)
self.optimizer = optim.SGD(
self.yolov4.parameters(),
lr=cfg.TRAIN["LR_INIT"],
momentum=cfg.TRAIN["MOMENTUM"],
weight_decay=cfg.TRAIN["WEIGHT_DECAY"],
)
self.criterion = YoloV4Loss(
anchors=cfg.MODEL["ANCHORS"],
strides=cfg.MODEL["STRIDES"],
iou_threshold_loss=cfg.TRAIN["IOU_THRESHOLD_LOSS"],
)
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),
)
if resume:
self.__load_resume_weights(weight_path)
