def __init__(self):
# ------------------------------- Init Network ---------------------------------#
self.config = TaskConfig()
# ------------------------------ Set Random Sed -------------------------------#
if self.config.set_seed:
seed_torch() # Fix random seed for try different Hyper-Parameter
save_log(self.config.log_file_path, "Set Random Sed")
# ------------------------------- Set Training Set ------------------------------#
self.dataset = DataSetCoco(split="train", config=self.config)
update_print_loss_interval(self.config, len(self.dataset))
self.train_loader = torch.utils.data.DataLoader(
self.dataset,
batch_size=self.config.batch_size,
shuffle=True,
num_workers=self.config.num_workers,
pin_memory=True)
# ------------------------------- Set Evaltor ----------------------------------#
# self.evaltor = Evaltor()
# ------------------------------ Init Network ----------------------------------#
if "dla" in self.config.model_name:
self.network = dladcn.getDlaDCN(34,
heads={
"hm": self.config.class_num,
"reg": 2,
"wh": 2
}).cuda(self.config.gpu_num[0])
# ------------------------------ Set Optimizer ----------------------------------#
# Use SGD
if self.config.which_optimizer == "sgd":
self.optimizer = torch.optim.SGD(self.network.parameters(),
lr=self.config.base_lr,
momentum=0.9,
weight_decay=1e-5)
# Use ADAM
elif self.config.which_optimizer == "adam":
self.optimizer = optim.Adam(self.network.parameters(),
lr=self.config.base_lr)
# ------------------------------- Init Loss -----------------------------------#
self.network_loss = CtdetLoss(self.config).cuda()
self.train_epoch_batchs = len(self.train_loader)
border_mode=cv2.BORDER_CONSTANT,
value=0,
p=0.2),
A.RandomCrop(height=480, width=640, p=0.2),
# A.RandomRotate90()
# A.RandomScale(always_apply=True, scale_limit=(0.5, 2.0))
A.Resize(height=480, width=640)
],
bbox_params=A.BboxParams(format='pascal_voc',
min_visibility=0.3,
label_fields=['category_ids']),
)
random.seed(7)
cocoRoot = "/data/yyh/2020/Simple-Centernet-with-Visualization-Pytorch_/dataset/code_testing_coco"
dataType = "val2017"
config = TaskConfig()
dataset = DataSetCoco("val", config)
for img_id in dataset.images:
imgInfo = dataset.coco.loadImgs(img_id)[0]
imPath = os.path.join(cocoRoot, dataType, imgInfo['file_name'])
print(imPath)
annIds = dataset.coco.getAnnIds(imgIds=imgInfo['id'])
im = cv2.imread(imPath)
origin_img = copy.deepcopy(im)
anns = dataset.coco.loadAnns(annIds)
bbox_list = []
label_name_list = []
for single_boox in anns:
label_name = dataset.coco.loadCats([single_boox["category_id"]
])[0]["name"]
# if (iteration % self.config.print_loss_interval) == self.config.print_loss_remainder or iteration == 0:
# one_iteration_time = time.time() - iteration_start_time
# train_info = '[Epoch: %d,Iteration: %d] loss:%.06f lr= %f time=%.02f s' % (
# epoch, iteration + 1, loss.item(), self.optimizer.param_groups[0]['lr'], one_iteration_time)
# self.config.logger.info(train_info)
# one_epoch_time = "One epoch spend {} minutes!!!!".format(
# (time.time() - epoch_start_time) / 60)
# self.config.logger.info(one_epoch_time)
# ------------------------------- Evaluate ----------------------------------#
self.net = self.evaluator.eval_val(self.net, epoch)
if self.config.best_acc >= config.model_ready_acc and epoch >= config.model_ready_epoch:
self.config.logger.info(
"Val Acc{} >= {}, Early Stop !!!! Best Epoch is {}".format(self.config.best_acc, config.model_ready_acc,
epoch))
break
# ------------------------------ Save Model ---------------------------------#
save_model(config)
if __name__ == "__main__":
# # efficientnet-b0
# config = TaskConfig("efficientnet-b0")
# tainer = Trainer(config)
# tainer.train()
# resnet50
config = TaskConfig("resnet50")
tainer = Trainer(config)
tainer.train()