def __init__(self, cfg_path):
with open(cfg_path, 'r') as rf:
self.cfg = yaml.safe_load(rf)
self.data_cfg = self.cfg['data']
self.model_cfg = self.cfg['model']
self.optim_cfg = self.cfg['optim']
self.hyper_params = self.cfg['hyper_params']
self.val_cfg = self.cfg['val']
print(self.data_cfg)
print(self.model_cfg)
print(self.optim_cfg)
print(self.hyper_params)
print(self.val_cfg)
os.environ['CUDA_VISIBLE_DEVICES'] = self.cfg['gpus']
dist.init_process_group(backend='nccl')
self.tdata = COCODataSets(
img_root=self.data_cfg['train_img_root'],
annotation_path=self.data_cfg['train_annotation_path'],
img_size=self.data_cfg['img_size'],
debug=self.data_cfg['debug'],
use_crowd=self.data_cfg['use_crowd'],
augments=True,
remove_blank=self.data_cfg['remove_blank'])
self.tloader = DataLoader(dataset=self.tdata,
batch_size=self.data_cfg['batch_size'],
num_workers=self.data_cfg['num_workers'],
collate_fn=self.tdata.collate_fn,
sampler=DistributedSampler(
dataset=self.tdata, shuffle=True))
self.vdata = COCODataSets(
img_root=self.data_cfg['val_img_root'],
annotation_path=self.data_cfg['val_annotation_path'],
img_size=self.data_cfg['img_size'],
debug=self.data_cfg['debug'],
use_crowd=self.data_cfg['use_crowd'],
augments=False,
remove_blank=False)
self.vloader = DataLoader(dataset=self.vdata,
batch_size=self.data_cfg['batch_size'],
num_workers=self.data_cfg['num_workers'],
collate_fn=self.vdata.collate_fn,
sampler=DistributedSampler(
dataset=self.vdata, shuffle=False))
print("train_data: ", len(self.tdata), " | ", "val_data: ",
len(self.vdata), " | ", "empty_data: ",
self.tdata.empty_images_len)
print("train_iter: ", len(self.tloader), " | ", "val_iter: ",
len(self.vloader))
model = RetinaNet(
num_cls=self.model_cfg['num_cls'],
anchor_sizes=self.model_cfg['anchor_sizes'],
strides=self.model_cfg['strides'],
backbone=self.model_cfg['backbone'],
)
if self.model_cfg.get("backbone_weight", None):
weights = torch.load(self.model_cfg['backbone_weight'])
model.load_backbone_weighs(weights)
self.best_map = 0.
self.best_map50 = 0.
optimizer = split_optimizer(model, self.optim_cfg)
local_rank = dist.get_rank()
self.local_rank = local_rank
self.device = torch.device("cuda", local_rank)
model.to(self.device)
model, optimizer = amp.initialize(model,
optimizer,
opt_level='O1',
verbosity=0)
if self.optim_cfg['sync_bn']:
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
self.model = nn.parallel.distributed.DistributedDataParallel(
model, device_ids=[local_rank], output_device=local_rank)
self.optimizer = optimizer
self.ema = ModelEMA(self.model)
beta = eval(self.hyper_params['beta']) if isinstance(self.hyper_params['beta'], str) \
else self.hyper_params['beta']
self.creterion = RetinaAnchorFreeLoss(
alpha=self.hyper_params['alpha'],
gamma=self.hyper_params['gamma'],
beta=beta,
top_k=self.hyper_params['top_k'],
box_iou_thresh=self.hyper_params['box_iou_thresh'],
box_reg_weight=self.hyper_params['box_reg_weight'])
self.lr_adjuster = WarmUpCosineDecayMultiStepLRAdjust(
init_lr=self.optim_cfg['lr'],
milestones=self.optim_cfg['milestones'],
warm_up_epoch=self.optim_cfg['warm_up_epoch'],
iter_per_epoch=len(self.tloader),
epochs=self.optim_cfg['epochs'],
cosine_weights=self.optim_cfg['cosine_weights'])
def __init__(self, cfg_path):
with open(cfg_path, 'r') as rf:
self.cfg = yaml.safe_load(rf)
self.data_cfg = self.cfg['data']
self.model_cfg = self.cfg['model']
self.optim_cfg = self.cfg['optim']
self.hyper_params = self.cfg['hyper_params']
self.val_cfg = self.cfg['val']
print(self.data_cfg)
print(self.model_cfg)
print(self.optim_cfg)
print(self.hyper_params)
print(self.val_cfg)
self.tdata = COCODataSets(
img_root=self.data_cfg['train_img_root'],
annotation_path=self.data_cfg['train_annotation_path'],
img_size=self.data_cfg['img_size'],
debug=self.data_cfg['debug'],
use_crowd=self.data_cfg['use_crowd'],
augments=True,
remove_blank=self.data_cfg['remove_blank'])
self.tloader = DataLoader(dataset=self.tdata,
batch_size=self.data_cfg['batch_size'],
num_workers=self.data_cfg['num_workers'],
collate_fn=self.tdata.collate_fn,
shuffle=True)
self.vdata = COCODataSets(
img_root=self.data_cfg['val_img_root'],
annotation_path=self.data_cfg['val_annotation_path'],
img_size=self.data_cfg['img_size'],
debug=self.data_cfg['debug'],
use_crowd=self.data_cfg['use_crowd'],
augments=False,
remove_blank=False)
self.vloader = DataLoader(dataset=self.vdata,
batch_size=self.data_cfg['batch_size'],
num_workers=self.data_cfg['num_workers'],
collate_fn=self.vdata.collate_fn,
shuffle=False)
print("train_data: ", len(self.tdata), " | ", "val_data: ",
len(self.vdata), " | ", "empty_data: ",
self.tdata.empty_images_len)
print("train_iter: ", len(self.tloader), " | ", "val_iter: ",
len(self.vloader))
model = RetinaNet(
num_cls=self.model_cfg['num_cls'],
anchor_sizes=self.model_cfg['anchor_sizes'],
strides=self.model_cfg['strides'],
backbone=self.model_cfg['backbone'],
)
if self.model_cfg.get("backbone_weight", None):
weights = torch.load(self.model_cfg['backbone_weight'])
model.load_backbone_weighs(weights)
self.best_map = 0.
self.best_map50 = 0.
optimizer = split_optimizer(model, self.optim_cfg)
self.local_rank = 0
self.device = torch.device("cuda:0")
model.to(self.device)
self.model = model
self.optimizer = optimizer
self.ema = ModelEMA(self.model)
beta = eval(self.hyper_params['beta']) if isinstance(self.hyper_params['beta'], str) \
else self.hyper_params['beta']
self.creterion = RetinaAnchorFreeLoss(
alpha=self.hyper_params['alpha'],
gamma=self.hyper_params['gamma'],
beta=beta,
top_k=self.hyper_params['top_k'],
box_iou_thresh=self.hyper_params['box_iou_thresh'],
box_reg_weight=self.hyper_params['box_reg_weight'])
self.lr_adjuster = WarmUpCosineDecayMultiStepLRAdjust(
init_lr=self.optim_cfg['lr'],
milestones=self.optim_cfg['milestones'],
warm_up_epoch=self.optim_cfg['warm_up_epoch'],
iter_per_epoch=len(self.tloader),
epochs=self.optim_cfg['epochs'],
cosine_weights=self.optim_cfg['cosine_weights'])