dataset = ICDARDataset(config.icdar17_mlt_img_dir,
config.icdar17_mlt_gt_dir)
dataloader = DataLoader(dataset,
batch_size=1,
shuffle=True,
num_workers=config.num_workers)
model = CTPN_Model()
model.to(device)
if os.path.exists(checkpoints_weight):
print('using pretrained weight: {}'.format(checkpoints_weight))
cc = torch.load(checkpoints_weight, map_location=device)
model.load_state_dict(cc['model_state_dict'])
resume_epoch = cc['epoch']
else:
model.apply(weights_init)
params_to_uodate = model.parameters()
optimizer = optim.SGD(params_to_uodate, lr=lr, momentum=0.9)
critetion_cls = RPN_CLS_Loss(device)
critetion_regr = RPN_REGR_Loss(device)
best_loss_cls = 100
best_loss_regr = 100
best_loss = 100
best_model = None
epochs += resume_epoch
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
for epoch in range(resume_epoch + 1, epochs):
config.icdar19_mlt_gt_dir)
dataloader = DataLoader(dataset,
batch_size=1,
shuffle=True,
num_workers=config.num_workers)
model = CTPN_Model()
model.to(device)
if os.path.exists(checkpoints_weight):
print('using pretrained weight: {}'.format(checkpoints_weight))
cc = torch.load(checkpoints_weight, map_location=device)
model.load_state_dict(cc['model_state_dict'])
resume_epoch = cc['epoch']
else:
model.apply(
weights_init
) ## 函数-Module.apply(fn)：会递归地搜索网络内的所有module并把参数表示的函数应用到所有的module上。
params_to_update = model.parameters()
optimizer = optim.SGD(params_to_update, lr=lr, momentum=0.9)
critetion_cls = RPN_CLS_Loss(device)
critetion_regr = RPN_REGR_Loss(device)
best_loss_cls = 100
best_loss_regr = 100
best_loss = 100
best_model = None
epochs += resume_epoch
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)