def __init__(self,
model,
weight_path=None,
gpu_id=0,
img_size=544,
visiual=None,
eval=False):
self.model = model
self.img_size = img_size
self.__num_class = cfg.DATA["NUM"]
self.__conf_threshold = cfg.TEST["CONF_THRESH"]
self.__nms_threshold = cfg.TEST["NMS_THRESH"]
self.__device = gpu.select_device(gpu_id)
self.__multi_scale_test = cfg.TEST["MULTI_SCALE_TEST"]
self.__flip_test = cfg.TEST["FLIP_TEST"]
self.__visiual = visiual
self.__eval = eval
self.__classes = cfg.DATA["CLASSES"]
if self.model == 's':
self.__model = Yolov3_S().to(self.__device)
else:
self.__model = Yolov3().to(self.__device)
self.__load_model_weights(weight_path)
self.__evalter = Evaluator(self.__model, visiual=False)
def __init__(
self,
cfg,
use_weight_path,
img_dir,
gpu_id=0,
img_size=488,
annotation_dir='',
):
self.cfg = cfg
self.img_size = img_size
self.__num_class = cfg.DATA["NUM"]
self.__conf_threshold = cfg.TEST["CONF_THRESH"]
self.__nms_threshold = cfg.TEST["NMS_THRESH"]
self.__visual_threshold = cfg.TEST["VISUAL_THRESH"]
self.__device = gpu.select_device(gpu_id)
self.__multi_scale_test = cfg.TEST["MULTI_SCALE_TEST"]
self.__flip_test = cfg.TEST["FLIP_TEST"]
self.__img_dir = img_dir
self.__annotation_dir = annotation_dir
self.__classes = cfg.DATA["CLASSES"]
self.__model = Yolov3(cfg).to(self.__device)
self.__load_model_weights(use_weight_path)
self.__evalter = YoloEvaluator(self.__model,
cfg,
img_dir,
mode="test",
anno_dir=annotation_dir)
self.__visual_imgs = 0
def __init__(self, cfg, train_dir, valid_dir, weight_path, resume, gpu_id):
init_seeds(0)
self.cfg = cfg
self.train_dir = train_dir
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.today = datetime.datetime.today().strftime('%y%m%d')
num_weights = len(glob.glob(f'./weight/{self.today}_*'))
self.save_weight_dir = f'./weight/{self.today}_{num_weights+1}'
os.mkdir(self.save_weight_dir)
self.multi_scale_train = cfg.TRAIN["MULTI_SCALE_TRAIN"]
# self.train_dataset = VocDataset(cfg=cfg, anno_file_type="train", img_size=cfg.TRAIN["TRAIN_IMG_SIZE"])
self.train_dataset = YoloDataset(cfg=cfg, data_dir=self.train_dir, 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)
self.yolov3 = Yolov3(cfg=cfg).to(self.device)
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 = 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))
self.valid_dir = valid_dir
self.img_valid_dir = os.path.join(self.valid_dir, 'images')
self.anno_valid_dir = os.path.join(self.valid_dir, 'labels')
self.valid_dataset = YoloDataset(cfg=cfg, data_dir=self.valid_dir, img_size=cfg.EVAL["TEST_IMG_SIZE"])
self.valid_dataloader = DataLoader(self.valid_dataset,
batch_size=cfg.EVAL["BATCH_SIZE"],
num_workers=cfg.EVAL["NUMBER_WORKERS"],
shuffle=False)
self.evaluator = YoloEvaluator(self.yolov3, self.cfg, self.img_valid_dir, "eval", anno_dir=self.anno_valid_dir)
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))
# loss classes
loss_cls = label_obj_mask * BCE(input=p_cls,
target=label_cls) * label_mix
loss_giou = (torch.sum(loss_giou)) / batch_size
loss_conf = (torch.sum(loss_conf)) / batch_size
loss_cls = (torch.sum(loss_cls)) / batch_size
loss = loss_giou + loss_conf + loss_cls
return loss, loss_giou, loss_conf, loss_cls
if __name__ == "__main__":
from model.yolov3 import Yolov3
net = Yolov3()
p, p_d = net(torch.rand(3, 3, 416, 416))
label_sbbox = torch.rand(3, 52, 52, 3, 26)
label_mbbox = torch.rand(3, 26, 26, 3, 26)
label_lbbox = torch.rand(3, 13, 13, 3, 26)
sbboxes = torch.rand(3, 150, 4)
mbboxes = torch.rand(3, 150, 4)
lbboxes = torch.rand(3, 150, 4)
loss, loss_xywh, loss_conf, loss_cls = YoloV3Loss(
cfg.MODEL["ANCHORS"], cfg.MODEL["STRIDES"])(p, p_d, label_sbbox,
label_mbbox, label_lbbox,
sbboxes, mbboxes, lbboxes)
print(loss)
# train_state_dict = torch.load('./weight/repA1g4_seed1_20each/best.pt')
# train_model = Yolov3(cfg=cfg, deploy=False)
# train_model.load_state_dict(train_state_dict, strict=False)
# deploy_model = Yolov3(cfg=cfg, deploy=True)
# whole_model_convert(train_model, deploy_model)
# deploy_state_dict = deploy_model.state_dict()
# torch.save(deploy_state_dict, './weight/repA1g4_seed1_20each/best_deploy.pt')
################################# test eqivalent ##################################################
train_state_dict = torch.load('./weight/repA1g4_seed1_20each/best.pt')
train_model = Yolov3(cfg=cfg, deploy=False)
train_model.load_state_dict(train_state_dict, strict=False)
train_model.cuda()
deploy_state_dict = torch.load('./weight/repA1g4_seed1_20each/best_deploy.pt')
deploy_model = Yolov3(cfg=cfg, deploy=True)
deploy_model.load_state_dict(deploy_state_dict, strict=False)
deploy_model.cuda()
train_model.eval()
deploy_model.eval()
x = torch.randn(1, 3, 224, 224).cuda()
with torch.no_grad():
train_p, train_p_d = train_model(x)