def __init__(self, arg=None, vis=None, id=None):
super().__init__()
self.arg = ARG() if arg is None else arg
self.vis = VIS(save_dir=arg.save.tensorboard) if vis is None else vis
self.id = 0 if id is None else id
self.global_step = 0
self.hyper_parameter = dict()
self.hyper_parameter['lr'] = 10**-1.6
self.dataset = Hand_Dataset(
self.arg,
self.arg.data.dataset.train.root,
im_size=self.arg.model.net.im_size,
relative_path_txt=self.arg.data.dataset.train.txt_path)
self.dataloader = DataLoader(
self.dataset,
batch_size=self.arg.data.dataloader.train.batch_size,
shuffle=self.arg.data.dataloader.train.shuffle,
drop_last=self.arg.data.dataloader.train.drop_last)
# print('train dataloader len = ', len(self.dataloader))
self.net = YoloV3_Micro(class_num=self.arg.model.net.num_class,
anchors_num=self.arg.model.net.num_anchor)
input_to_net = torch.rand(1, 3, self.arg.model.net.im_size[0],
self.arg.model.net.im_size[1])
# self.net.forward(input_to_net)
self.vis.model(self.net, input_to_net)
if self.arg.model.weight_path is not None:
self.net.load_state_dict(
one_card_model(torch.load(self.arg.model.weight_path)))
# self.net = torch.nn.DataParallel(self.net)
self.loss_func = YOLOLoss()
# loss_func = torch.nn.DataParallel(loss_func)
self.optimizer = torch.optim.Adam(self.net.parameters(),
lr=self.hyper_parameter['lr'])
# self.lr_decay = ExponentialDecay(0.1, 10000, 1e-4)
self.bbox_format = list()
self.bbox_format.append(
Bbox_Format(self.arg.model.net.im_size,
self.arg.model.net.feature_size[0],
self.arg.model.flt_anchor,
self.arg.model.net.num_class,
self.arg.model.mask_iou_threshold))
self.bbox_format.append(
Bbox_Format(self.arg.model.net.im_size,
self.arg.model.net.feature_size[1],
self.arg.model.flt_anchor,
self.arg.model.net.num_class,
self.arg.model.mask_iou_threshold))
self.valider = Valider(self.arg)
self.tester = Tester(self.arg)
def __init__(self, arg):
super().__init__()
self.arg = arg
self.dataset = Hand_Dataset(arg, arg.data.dataset.test.root, im_size=arg.model.net.im_size, mode='test')
self.dataloader = DataLoader(self.dataset, batch_size=1, shuffle=False)
self.bbox_format = list()
self.bbox_format.append(Bbox_Format(arg.model.net.im_size, arg.model.net.feature_size[0], arg.model.flt_anchor,
arg.model.net.num_class, arg.model.mask_iou_threshold))
self.bbox_format.append(Bbox_Format(arg.model.net.im_size, arg.model.net.feature_size[1], arg.model.flt_anchor,
arg.model.net.num_class, arg.model.mask_iou_threshold))
def main():
arg = ARG()
prepare_save_dirs(arg)
vis = VIS(save_dir=arg.save.tensorboard)
# =======================================================
train_dataset = Hand_Dataset(arg, arg.data.dataset.train.root, im_size=arg.model.net.im_size, relative_path_txt=arg.data.dataset.train.txt_path)
train_loader = DataLoader(train_dataset, batch_size=arg.data.dataloader.train.batch_size, shuffle=arg.data.dataloader.train.shuffle, drop_last=arg.data.dataloader.train.drop_last)
vis.iteration_per_epoch = len(train_loader)
valid_dataset = Hand_Dataset(arg, arg.data.dataset.valid.root, im_size=arg.model.net.im_size, mode='valid')
valid_loader = DataLoader(valid_dataset, batch_size=1, shuffle=False)
test_dataset = Hand_Dataset(arg, arg.data.dataset.test.root, im_size=arg.model.net.im_size, mode='test')
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
net = YoloV3_Micro(class_num=arg.model.net.num_class, anchors_num=arg.model.net.num_anchor)
input_to_net = torch.rand(1, 3, arg.model.net.im_size[0], arg.model.net.im_size[1])
net.forward(input_to_net)
vis.model(net, input_to_net)
net = net.to(arg.train.device)
if arg.model.weight_path is not None:
net.load_state_dict(one_card_model(torch.load(arg.model.weight_path)))
net = torch.nn.DataParallel(net)
loss_func = YOLOLoss(arg.train.device)
loss_func = loss_func.to(arg.train.device)
loss_func = torch.nn.DataParallel(loss_func)
optimizer = torch.optim.Adam(net.parameters(), lr=arg.optim.lr, weight_decay=arg.optim.weight_decay)
lr_decay = LrSeeker(1e-9, 1)
for epoch in range(1, arg.train.epochs + 1):
vis.epoch_counter = epoch-1
train(arg, train_loader, valid_loader, net, loss_func, optimizer, lr_decay, epoch, vis)
def main():
arg = ARG()
prepare_save_dirs(arg)
vis = VIS(save_dir=arg.save.tensorboard)
# =======================================================
hand_dataset = Hand_Dataset(arg, arg.data.dataset.train.root,
arg.data.dataset.train.relative_txt_path,
arg.model.net.im_size)
train_loader = DataLoader(hand_dataset, batch_size=16, shuffle=False)
vis.iteration_per_epoch = len(train_loader)
net = YoloV3_Micro(class_num=arg.model.net.num_class,
anchors_num=arg.model.net.num_anchor)
input_to_net = torch.rand(1, 3, arg.model.net.im_size[0],
arg.model.net.im_size[1])
net.forward(input_to_net)
vis.model(net, input_to_net)
net = net.to(arg.train.device)
if arg.model.weight_path is not None:
net.load_state_dict(one_card_model(torch.load(arg.model.weight_path)))
net = torch.nn.DataParallel(net)
loss_func = YOLOLoss(arg.train.device)
loss_func = loss_func.to(arg.train.device)
loss_func = torch.nn.DataParallel(loss_func)
optimizer = torch.optim.Adam(net.parameters(),
lr=arg.optim.lr,
weight_decay=arg.optim.weight_decay)
for epoch in range(1, arg.train.epochs + 1):
vis.epoch_counter = epoch
train(arg, train_loader, net, loss_func, optimizer, epoch, vis)
# if epoch % arg.train.valid_interval == 0:
# valid(arg, valid_loader, model, epoch, vis)
if epoch % arg.train.save_model_epoch_interval == 0:
torch.save(
net.state_dict(),
os.path.join(arg.save.model,
'micro_yolo_epoch{}_weight'.format(epoch)))
cv2.imshow('', np_im)
cv2.waitKey()
os.makedirs('image_out', exist_ok=True)
cv2.imwrite('image_out/' + str(image_counter) + '.jpg', np_im)
image_counter += 1
rec, prec, ap = det_eval(predict_list, target_dict, 0, 0.5)
return ap
if __name__ == '__main__':
from gear_config.yolo_config.body_micro_yolo_aug1 import ARG
arg = ARG()
test_dataset = Hand_Dataset(arg,
arg.data.dataset.test.root,
im_size=arg.model.net.im_size,
mode='test')
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
net = YoloV3_Micro(class_num=arg.model.net.num_class,
anchors_num=arg.model.net.num_anchor)
net = net.to(arg.test.device)
if arg.model.weight_path is not None:
net.load_state_dict(one_card_model(torch.load(arg.model.weight_path)))
tst(arg, test_loader, net, write_image=True)
def main():
arg = ARG()
prepare_save_dirs(arg)
vis = VIS(save_dir=arg.save.tensorboard)
# =======================================================
train_dataset = Hand_Dataset(
arg,
arg.data.dataset.train.root,
im_size=arg.model.net.im_size,
relative_path_txt=arg.data.dataset.train.txt_path)
train_loader = DataLoader(train_dataset,
batch_size=arg.data.dataloader.train.batch_size,
shuffle=arg.data.dataloader.train.shuffle,
drop_last=arg.data.dataloader.train.drop_last)
vis.iteration_per_epoch = len(train_loader)
valid_dataset = Hand_Dataset(arg,
arg.data.dataset.valid.root,
im_size=arg.model.net.im_size,
mode='valid')
valid_loader = DataLoader(valid_dataset, batch_size=1, shuffle=False)
test_dataset = Hand_Dataset(arg,
arg.data.dataset.test.root,
im_size=arg.model.net.im_size,
mode='test')
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
net = YoloV3_Micro(class_num=arg.model.net.num_class,
anchors_num=arg.model.net.num_anchor)
input_to_net = torch.rand(1, 3, arg.model.net.im_size[0],
arg.model.net.im_size[1])
net.forward(input_to_net)
vis.model(net, input_to_net)
net = net.to(arg.train.device)
if arg.model.weight_path is not None:
net.load_state_dict(one_card_model(torch.load(arg.model.weight_path)))
net = torch.nn.DataParallel(net)
loss_func = YOLOLoss(arg.train.device)
loss_func = loss_func.to(arg.train.device)
loss_func = torch.nn.DataParallel(loss_func)
optimizer = torch.optim.Adam(net.parameters(), lr=arg.optim.lr)
lr_decay = ExponentialDecay(0.1, 10000, 1e-4)
for epoch in range(1, arg.train.epochs + 1):
vis.epoch_counter = epoch - 1
train(arg, train_loader, valid_loader, net, loss_func, optimizer,
lr_decay, epoch, vis)
if epoch % arg.train.save_model_epoch_interval == 0:
if arg.save.model is not None:
torch.save(
net.state_dict(),
os.path.join(
arg.save.model,
'body_micro_yolo_epoch{}_weight'.format(epoch)))
if epoch % arg.train.test_epoch_interval == 0:
ap = tst(arg, test_loader, net)
vis.line('test/ap', ap)