def main(model_path, backbone, scale, path, save_path, gpu_id):
if os.path.exists(save_path):
shutil.rmtree(save_path, ignore_errors=True)
if not os.path.exists(save_path):
os.makedirs(save_path)
save_img_folder = os.path.join(save_path, 'img')
if not os.path.exists(save_img_folder):
os.makedirs(save_img_folder)
save_txt_folder = os.path.join(save_path, 'result')
if not os.path.exists(save_txt_folder):
os.makedirs(save_txt_folder)
img_paths = [os.path.join(path, x) for x in os.listdir(path)]
net = PSENet(backbone=backbone, pretrained=False, result_num=config.n)
model = Pytorch_model(model_path, net=net, scale=scale, gpu_id=gpu_id)
total_frame = 0.0
total_time = 0.0
for img_path in tqdm(img_paths):
img_name = os.path.basename(img_path).split('.')[0]
save_name = os.path.join(save_txt_folder, 'res_' + img_name + '.txt')
_, boxes_list, t = model.predict(img_path)
total_frame += 1
total_time += t
img = draw_bbox(img_path, boxes_list, color=(0, 0, 255), thickness=7)
cv2.imwrite(os.path.join(save_img_folder, '{}.jpg'.format(img_name)),
img)
np.savetxt(save_name,
boxes_list.reshape(-1, 8),
delimiter=',',
fmt='%d')
print('fps:{}'.format(total_frame / total_time))
return save_txt_folder
def main(model_path, backbone, scale, path, save_path, gpu_id):
device = torch.device("cuda:" + str(gpu_id))
logger = setup_logger(os.path.join(config.output_dir, 'test_log'))
logger.info(config.print())
if os.path.exists(save_path):
shutil.rmtree(save_path, ignore_errors=True)
if not os.path.exists(save_path):
os.makedirs(save_path)
save_img_folder = os.path.join(save_path, 'img')
if not os.path.exists(save_img_folder):
os.makedirs(save_img_folder)
save_txt_folder = os.path.join(save_path, 'result')
if not os.path.exists(save_txt_folder):
os.makedirs(save_txt_folder)
img_paths = [os.path.join(path, x) for x in os.listdir(path)]
net = PSENet(backbone=backbone,
pretrained=config.pretrained,
result_num=config.n)
model = Pytorch_model(model_path, net=net, scale=scale, gpu_id=gpu_id)
num_gpus = torch.cuda.device_count()
if num_gpus > 1:
model = nn.DataParallel(model)
model = model.to(device)
model = nn.DataParallel(model)
recall, precision, f1 = merge_eval(model=model,
save_path=os.path.join(
config.output_dir, 'output'),
test_path=config.testroot,
device=device,
base_path=config.base_path,
use_sub=config.use_sub)
logger.info('test: recall: {:.6f}, precision: {:.6f}, f1: {:.6f}'.format(
recall, precision, f1))
# total_frame = 0.0
# total_time = 0.0
# for img_path in tqdm(img_paths):
# img_name = os.path.basename(img_path).split('.')[0]
# save_name = os.path.join(save_txt_folder, 'res_' + img_name + '.txt')
# _, boxes_list, t = model.predict(img_path)
# total_frame += 1
# total_time += t
# # img = draw_bbox(img_path, boxes_list, color=(0, 0, 255))
# # cv2.imwrite(os.path.join(save_img_folder, '{}.jpg'.format(img_name)), img)
# np.savetxt(save_name, boxes_list.reshape(-1, 8), delimiter=',', fmt='%d')
# print('fps:{}'.format(total_frame / total_time))
return save_txt_folder
def main():
if config.output_dir is None:
config.output_dir = 'output'
if config.restart_training:
shutil.rmtree(config.output_dir, ignore_errors=True)
if not os.path.exists(config.output_dir):
os.makedirs(config.output_dir)
logger = setup_logger(os.path.join(config.output_dir, 'train_log'))
logger.info(config.print())
torch.manual_seed(config.seed) # 为CPU设置随机种子
if config.gpu_id is not None and torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
logger.info('train with gpu {} and pytorch {}'.format(
config.gpu_id, torch.__version__))
device = torch.device("cuda:0")
torch.cuda.manual_seed(config.seed) # 为当前GPU设置随机种子
torch.cuda.manual_seed_all(config.seed) # 为所有GPU设置随机种子
else:
logger.info('train with cpu and pytorch {}'.format(torch.__version__))
device = torch.device("cpu")
train_data = TibetanDataset(config.json_path,
data_shape=config.data_shape,
n=config.n,
m=config.m,
transform=transforms.ToTensor(),
base_path=config.base_path)
train_loader = Data.DataLoader(dataset=train_data,
batch_size=config.train_batch_size,
shuffle=True,
num_workers=int(config.workers))
writer = SummaryWriter(config.output_dir)
model = PSENet(backbone=config.backbone,
pretrained=config.pretrained,
result_num=config.n,
scale=config.scale)
if not config.pretrained and not config.restart_training:
model.apply(weights_init)
num_gpus = torch.cuda.device_count()
if num_gpus > 1:
model = nn.DataParallel(model)
model = model.to(device)
# dummy_input = torch.autograd.Variable(torch.Tensor(1, 3, 600, 800).to(device))
# writer.add_graph(models=models, input_to_model=dummy_input)
criterion = PSELoss(Lambda=config.Lambda,
ratio=config.OHEM_ratio,
reduction='mean')
# optimizer = torch.optim.SGD(models.parameters(), lr=config.lr, momentum=0.99)
optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
if config.checkpoint != '' and not config.restart_training:
start_epoch = load_checkpoint(config.checkpoint, model, logger, device,
optimizer)
start_epoch += 1
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
config.lr_decay_step,
gamma=config.lr_gamma,
last_epoch=start_epoch)
else:
start_epoch = config.start_epoch
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
config.lr_decay_step,
gamma=config.lr_gamma)
all_step = len(train_loader)
logger.info('train dataset has {} samples,{} in dataloader'.format(
train_data.__len__(), all_step))
epoch = 0
best_model = {'recall': 0, 'precision': 0, 'f1': 0, 'models': ''}
try:
for epoch in range(start_epoch, config.epochs):
start = time.time()
train_loss, lr = train_epoch(model, optimizer, scheduler,
train_loader, device, criterion,
epoch, all_step, writer, logger)
logger.info(
'[{}/{}], train_loss: {:.4f}, time: {:.4f}, lr: {}'.format(
epoch, config.epochs, train_loss,
time.time() - start, lr))
# net_save_path = '{}/PSENet_{}_loss{:.6f}.pth'.format(config.output_dir, epoch,
# train_loss)
# save_checkpoint(net_save_path, models, optimizer, epoch, logger)
if (0.3 < train_loss < 0.4 and epoch % 4 == 0) or train_loss < 0.3:
recall, precision, f1 = merge_eval(model=model,
save_path=os.path.join(
config.output_dir,
'output'),
test_path=config.testroot,
device=device,
base_path=config.base_path,
use_sub=config.use_sub)
logger.info(
'test: recall: {:.6f}, precision: {:.6f}, f1: {:.6f}'.
format(recall, precision, f1))
net_save_path = '{}/PSENet_{}_loss{:.6f}_r{:.6f}_p{:.6f}_f1{:.6f}.pth'.format(
config.output_dir, epoch, train_loss, recall, precision,
f1)
save_checkpoint(net_save_path, model, optimizer, epoch, logger)
if f1 > best_model['f1']:
best_path = glob.glob(config.output_dir + '/Best_*.pth')
for b_path in best_path:
if os.path.exists(b_path):
os.remove(b_path)
best_model['recall'] = recall
best_model['precision'] = precision
best_model['f1'] = f1
best_model['models'] = net_save_path
best_save_path = '{}/Best_{}_r{:.6f}_p{:.6f}_f1{:.6f}.pth'.format(
config.output_dir, epoch, recall, precision, f1)
if os.path.exists(net_save_path):
shutil.copyfile(net_save_path, best_save_path)
else:
save_checkpoint(best_save_path, model, optimizer,
epoch, logger)
pse_path = glob.glob(config.output_dir + '/PSENet_*.pth')
for p_path in pse_path:
if os.path.exists(p_path):
os.remove(p_path)
writer.add_scalar(tag='Test/recall',
scalar_value=recall,
global_step=epoch)
writer.add_scalar(tag='Test/precision',
scalar_value=precision,
global_step=epoch)
writer.add_scalar(tag='Test/f1',
scalar_value=f1,
global_step=epoch)
writer.close()
except KeyboardInterrupt:
save_checkpoint('{}/final.pth'.format(config.output_dir), model,
optimizer, epoch, logger)
finally:
if best_model['models']:
logger.info(best_model)
# model_path = 'output/psenet_icd2015_new_loss/final.pth'
image_root = 'Test Set/Image'
annotation_root = 'Test Set/Annotation'
Images = os.listdir(image_root)
Labels = os.listdir(annotation_root)
image_id = Images[np.random.randint(1, len(Images) + 1)].split('.jpg')[0]
img_path = image_root + os.sep + '{}.jpg'.format(image_id)
label_path = annotation_root + os.sep + '{}.txt'.format(image_id)
print('Predicting for image', img_path)
label = _get_annotation(label_path)
# 初始化网络
net = PSENet(backbone='resnet18', pretrained=False, result_num=config.n)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Pytorch_model(model_path, net=net, scale=1, device=device)
# for i in range(100):
# models.predict(img_path)
preds, boxes_list, t = model.predict(img_path)
print(boxes_list)
show_img(preds)
img = draw_bbox(img_path, boxes_list, color=(0, 0, 255))
cv2.imwrite('result.jpg', img)
# img = draw_bbox(img, label,color=(0,0,255))
show_img(img, color=True)
plt.show()
print('load label failed on {}'.format(label_path))
return np.array(boxes, dtype=np.float32)
if __name__ == '__main__':
from configs import config_tips as config
from models import PSENet
from utils.utils import show_img, draw_bbox
#model_path = 'outp ut/psenet_icd2015_resnet152_author_crop_adam_warm_up_myloss/best_r0.714011_p0.708214_f10.711100.pth'
#model_path = 'output/psenet_icd2015_new_loss/final.pth'
model_path = 'checkpoint/pse_epoch_30.pth'
#image_files = Path("/home/peizhao/data/icdar/2019/tips/test/img").rglob('*.jpg')
#image_files = Path("/home/peizhao/data/temp/doc").rglob("*.jpg")
image_files = Path("/home/peizhao/data/temp").rglob("*.png")
# image_files = Path("/home/peizhao/data/temp").rglob("*.jpg")
#image_files = Path("/home/peizhao/data/temp/test").rglob("*.jpg")
# 初始化网络
#net = PSENet(backbone='resnet152', pretrained=False, result_num=config.n)
net = PSENet(backbone='resnet50', pretrained=False, result_num=3)
model = Pytorch_model2(model_path, net=net, scale=1, gpu_id=0)
for item in image_files:
preds, boxes_list, t = model.predict(str(item))
img = draw_bbox(str(item), boxes_list, color=(0, 0, 255))
h, w = img.shape[:2]
scale = 640 / max(h, w)
img_size = cv2.resize(img, None, fx=scale, fy=scale)
cv2.imshow("result", img_size)
cv2.waitKey(0)
def main():
config.workspace = os.path.join(config.workspace_dir, config.exp_name)
if not os.path.exists(config.workspace):
os.makedirs(config.workspace)
logger = setup_logger(os.path.join(config.workspace, 'train_log'))
logger.info(config.pprint())
torch.manual_seed(config.seed) # 为CPU设置随机种子
torch.backends.cudnn.benchmark = True
logger.info('train with gpu {} and pytorch {}'.format(
config.gpu_id, torch.__version__))
device = torch.device("cuda:0")
torch.cuda.manual_seed(config.seed) # 为当前GPU设置随机种子
torch.cuda.manual_seed_all(config.seed) # 为所有GPU设置随机种子
train_data = TotalTextoader(config.train_data_dir,
config.train_gt_dir,
config.test_data_dir,
config.test_gt_dir,
split='train',
is_transform=True,
img_size=config.data_shape,
kernel_num=config.kernel_num,
min_scale=config.min_scale)
train_loader = Data.DataLoader(dataset=train_data,
batch_size=config.train_batch_size,
shuffle=True,
num_workers=int(config.workers))
model = PSENet(backbone=config.backbone,
pretrained=config.pretrained,
result_num=config.kernel_num,
scale=config.scale)
if not config.pretrained and not config.restart_training:
model.apply(weights_init)
num_gpus = torch.cuda.device_count()
# if num_gpus > 1:
model = nn.DataParallel(model)
model = model.to(device)
criterion = dice_loss
optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
if config.checkpoint != '' and config.restart_training == True:
start_epoch = load_checkpoint(config.checkpoint, model, logger, device,
optimizer)
start_epoch += 1
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
config.lr_decay_step,
gamma=config.lr_gamma,
last_epoch=start_epoch)
logger.info('resume from {}, epoch={}'.format(config.checkpoint,
start_epoch))
else:
start_epoch = 1
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
config.lr_decay_step,
gamma=config.lr_gamma)
all_step = len(train_loader)
logger.info('train dataset has {} samples,{} iters in dataloader'.format(
train_data.__len__(), all_step))
for epoch in range(start_epoch, config.epochs + 1):
start = time.time()
train_loss, lr = train_epoch(model, optimizer, scheduler, train_loader,
device, criterion, epoch, all_step,
logger)
logger.info('[{}/{}], train_loss: {:.4f}, time: {:.4f}, lr: {}'.format(
epoch, config.epochs, train_loss,
time.time() - start, lr))
if epoch % config.save_interval == 0:
save_path = '{}/epoch_{}.pth'.format(config.workspace, epoch)
latest_path = '{}/latest.pth'.format(config.workspace)
save_checkpoint(save_path, model, optimizer, epoch, logger)
save_checkpoint(latest_path, model, optimizer, epoch, logger)
def main():
if config.output_dir is None:
config.output_dir = 'output'
if config.restart_training:
shutil.rmtree(config.output_dir, ignore_errors=True)
if not os.path.exists(config.output_dir):
os.makedirs(config.output_dir)
logger = setup_logger(os.path.join(config.output_dir, 'train_log'))
logger.info(config.print())
torch.manual_seed(config.seed) # 为CPU设置随机种子
if config.gpu_id is not None and torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
logger.info('train with gpu {} and pytorch {}'.format(
config.gpu_id, torch.__version__))
device = torch.device("cuda:0")
torch.cuda.manual_seed(config.seed) # 为当前GPU设置随机种子
torch.cuda.manual_seed_all(config.seed) # 为所有GPU设置随机种子
else:
logger.info('train with cpu and pytorch {}'.format(torch.__version__))
device = torch.device("cpu")
train_data = MyDataset(args.train_dir,
data_shape=config.data_shape,
n=config.n,
m=config.m,
transform=transforms.ToTensor())
train_loader = Data.DataLoader(dataset=train_data,
batch_size=args.batch_size,
shuffle=True,
num_workers=int(config.workers))
writer = SummaryWriter(config.output_dir)
model = PSENet(backbone=config.backbone,
pretrained=config.pretrained,
result_num=config.n,
scale=config.scale)
if not config.pretrained and not config.restart_training:
model.apply(weights_init)
if args.resume_model:
resume_model(model, args.resume_model)
num_gpus = torch.cuda.device_count()
if num_gpus > 1:
model = nn.DataParallel(model)
model = model.to(device)
# dummy_input = torch.autograd.Variable(torch.Tensor(1, 3, 600, 800).to(device))
# writer.add_graph(models=models, input_to_model=dummy_input)
criterion = PSELoss(Lambda=config.Lambda,
ratio=config.OHEM_ratio,
reduction='mean')
# optimizer = torch.optim.SGD(models.parameters(), lr=config.lr, momentum=0.99)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
if config.checkpoint != '' and not config.restart_training:
start_epoch = load_checkpoint(config.checkpoint, model, logger, device,
optimizer)
start_epoch += 1
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
config.lr_decay_step,
gamma=config.lr_gamma,
last_epoch=start_epoch)
else:
start_epoch = config.start_epoch
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
config.lr_decay_step,
gamma=config.lr_gamma)
all_step = len(train_loader)
logger.info('train dataset has {} samples,{} in dataloader'.format(
train_data.__len__(), all_step))
epoch = 0
best_model = {'recall': 0, 'precision': 0, 'f1': 0, 'models': ''}
try:
for epoch in range(start_epoch, args.epochs):
start = time.time()
train_loss, lr = train_epoch(model, optimizer, scheduler,
train_loader, device, criterion,
epoch, all_step, writer, logger)
logger.info(
'[{}/{}], train_loss: {:.4f}, time: {:.4f}, lr: {}'.format(
epoch, config.epochs, train_loss,
time.time() - start, lr))
if epoch % args.save_per_epoch == 0:
save_model(model, epoch)
writer.close()
except KeyboardInterrupt:
save_checkpoint('{}/final.pth'.format(config.output_dir), model,
optimizer, epoch, logger)
finally:
if best_model['models']:
logger.info(best_model)
from utils.utils import show_img, draw_bbox
os.environ['CUDA_VISIBLE_DEVICES'] = str('0')
model_path = '/home/gp/work/project/PSENet.pytorch.mobilenetv3-small/output/psenet_jizhuangxiang_MobileNetV3_Small_1gpu_author_crop_adam_MultiStepLR_authorloss/PSENet_10_loss0.283093.pth'
# model_path = 'output/psenet_icd2015_new_loss/final.pth'
#img_id = 10
#img_path = '/data2/dataset/ICD15/test/img/img_{}.jpg'.format(img_id)
#label_path = '/data2/dataset/ICD15/test/gt/gt_img_{}.txt'.format(img_id)
img_path = '/home/gp/work/project/PSENet.pytorch/data/jizhuangxiang/train/img/3.jpg'
label_path = '/home/gp/work/project/PSENet.pytorch/data/jizhuangxiang/train/gt/gt_3.txt'
label = _get_annotation(label_path)
# 初始化网络
net = PSENet(backbone='MobileNetV3_Small',
pretrained=False,
result_num=config.n)
model = Pytorch_model(model_path, net=net, scale=1, gpu_id=0)
# for i in range(100):
# models.predict(img_path)
preds, boxes_list, t = model.predict(img_path)
print(boxes_list)
show_img(preds)
img = draw_bbox(img_path, boxes_list, color=(0, 0, 255))
cv2.imwrite('result.jpg', img)
# img = draw_bbox(img, label,color=(0,0,255))
show_img(img, color=True)
plt.show()
shutil.rmtree(fid_path, ignore_errors=True)
precision, recall, hmean = evl_totaltext(save_path, gt_path, fid_path)
# f_score_new = getresult(save_path,config.gt_name)
return precision, recall, hmean
if __name__ == "__main__":
config.workspace = os.path.join(config.workspace_dir, config.exp_name)
logger = setup_logger(os.path.join(config.workspace, 'test_log'))
logger.info(config.print())
# best_save_path = '{}/Best_model_0.632154.pth'.format(config.workspace)
best_save_path = "/data/glusterfs_cv_04/11121171/CVPR_Text/PSENet_file/Total_Text/Best_model_0.787389.pth"
# writer = SummaryWriter(config.output_dir)
model = PSENet(backbone=config.backbone,
pretrained=config.pretrained,
result_num=config.kernel_num,
scale=config.scale)
num_gpus = torch.cuda.device_count()
device = torch.device("cuda:0")
# if num_gpus > 1:
model = nn.DataParallel(model)
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
start_epoch = load_checkpoint(best_save_path, model, logger, device,
optimizer)
for kern_size in range(3, 10):
kern_size_ = kern_size
print("kern_size:", kern_size_)
def main():
config.workspace = os.path.join(config.workspace_dir, config.exp_name)
if config.restart_training:
shutil.rmtree(config.workspace, ignore_errors=True)
if not os.path.exists(config.workspace):
os.makedirs(config.workspace)
shutil.rmtree(os.path.join(config.workspace, 'train_log'),
ignore_errors=True)
logger = setup_logger(os.path.join(config.workspace, 'train_log'))
logger.info(config.print())
torch.manual_seed(config.seed) # 为CPU设置随机种子
if config.gpu_id is not None and torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
logger.info('train with gpu {} and pytorch {}'.format(
config.gpu_id, torch.__version__))
device = torch.device("cuda:0")
torch.cuda.manual_seed(config.seed) # 为当前GPU设置随机种子
torch.cuda.manual_seed_all(config.seed) # 为所有GPU设置随机种子
else:
logger.info('train with cpu and pytorch {}'.format(torch.__version__))
device = torch.device("cpu")
train_data = ICDAR17(config.trainroot,
data_shape=config.data_shape,
n=config.kernel_num,
m=config.min_scale)
train_loader = Data.DataLoader(dataset=train_data,
batch_size=config.train_batch_size,
shuffle=True,
num_workers=int(config.workers))
# writer = SummaryWriter(config.output_dir)
model = PSENet(backbone=config.backbone,
pretrained=config.pretrained,
result_num=config.kernel_num,
scale=config.scale)
if not config.pretrained and not config.restart_training:
model.apply(weights_init)
num_gpus = torch.cuda.device_count()
if num_gpus > 1:
model = nn.DataParallel(model)
model = model.to(device)
criterion = PSELoss(Lambda=config.Lambda,
ratio=config.OHEM_ratio,
reduction='mean')
# optimizer = torch.optim.SGD(models.parameters(), lr=config.lr, momentum=0.99)
optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
if config.checkpoint != '' and not config.restart_training:
start_epoch = load_checkpoint(config.checkpoint, model, logger, device,
optimizer)
start_epoch += 1
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
config.lr_decay_step,
gamma=config.lr_gamma,
last_epoch=start_epoch)
else:
start_epoch = config.start_epoch
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
config.lr_decay_step,
gamma=config.lr_gamma)
all_step = len(train_loader)
logger.info('train dataset has {} samples,{} in dataloader'.format(
train_data.__len__(), all_step))
epoch = 0
f1 = 0
try:
for epoch in range(start_epoch, config.epochs):
start = time.time()
train_loss, lr = train_epoch(model, optimizer, scheduler,
train_loader, device, criterion,
epoch, all_step, logger)
logger.info(
'[{}/{}], train_loss: {:.4f}, time: {:.4f}, lr: {}'.format(
epoch, config.epochs, train_loss,
time.time() - start, lr))
save_path = '{}/epoch_model.pth'.format(config.workspace)
save_checkpoint(save_path, model, optimizer, epoch, logger)
if epoch >= 50 and epoch % 10 == 0:
f_score_new = eval(model,
os.path.join(config.workspace, 'output'),
config.testroot, device)
logger.info(' ---------------------------------------')
logger.info(' test: f_score : {:.6f}'.format(f_score_new))
logger.info(' ---------------------------------------')
if f_score_new > f1:
f1 = f_score_new
best_save_path = '{}/Best_model_{:.6f}.pth'.format(
config.workspace, f1)
save_checkpoint(best_save_path, model, optimizer, epoch,
logger)
# writer.add_scalar(tag='Test/recall', scalar_value=recall, global_step=epoch)
# writer.add_scalar(tag='Test/precision', scalar_value=precision, global_step=epoch)
# writer.add_scalar(tag='Test/f1', scalar_value=f1, global_step=epoch)
# writer.close()
except KeyboardInterrupt:
save_checkpoint('{}/final.pth'.format(config.workspace), model,
optimizer, epoch, logger)
def main():
config.workspace = os.path.join(config.workspace_dir, config.exp_name)
# if config.restart_training:
# shutil.rmtree(config.workspace, ignore_errors=True)
if not os.path.exists(config.workspace):
os.makedirs(config.workspace)
logger = setup_logger(os.path.join(config.workspace, 'train_log'))
logger.info(config.print())
torch.manual_seed(config.seed) # 为CPU设置随机种子
if config.gpu_id is not None and torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
logger.info('train with gpu {} and pytorch {}'.format(
config.gpu_id, torch.__version__))
device = torch.device("cuda:0")
torch.cuda.manual_seed(config.seed) # 为当前GPU设置随机种子
torch.cuda.manual_seed_all(config.seed) # 为所有GPU设置随机种子
else:
logger.info('train with cpu and pytorch {}'.format(torch.__version__))
device = torch.device("cpu")
train_data = ICDAR15(config.trainroot,
config.is_pseudo,
data_shape=config.data_shape,
n=config.kernel_num,
m=config.m)
train_loader = Data.DataLoader(dataset=train_data,
batch_size=config.train_batch_size,
shuffle=True,
num_workers=int(config.workers))
# writer = SummaryWriter(config.output_dir)
model = PSENet(backbone=config.backbone,
pretrained=config.pretrained,
result_num=config.kernel_num,
scale=config.scale)
if not config.pretrained and not config.restart_training:
model.apply(weights_init)
num_gpus = torch.cuda.device_count()
if num_gpus > 1:
model = nn.DataParallel(model)
model = model.to(device)
criterion = PSELoss(Lambda=config.Lambda,
ratio=config.OHEM_ratio,
reduction='mean')
optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
if config.checkpoint != '' and not config.restart_training:
start_epoch = load_checkpoint(config.checkpoint, model, logger, device,
optimizer)
start_epoch += 1
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
config.lr_decay_step,
gamma=config.lr_gamma,
last_epoch=start_epoch)
logger.info('resume from {}, epoch={}'.format(config.checkpoint,
start_epoch))
else:
start_epoch = config.start_epoch
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
config.lr_decay_step,
gamma=config.lr_gamma)
all_step = len(train_loader)
logger.info('train dataset has {} samples,{} in dataloader'.format(
train_data.__len__(), all_step))
epoch = 0
f1 = 0
for epoch in range(start_epoch, config.epochs):
start = time.time()
train_loss, lr = train_epoch(model, optimizer, scheduler, train_loader,
device, criterion, epoch, all_step,
logger)
logger.info('[{}/{}], train_loss: {:.4f}, time: {:.4f}, lr: {}'.format(
epoch, config.epochs, train_loss,
time.time() - start, lr))
if epoch % config.save_interval == 0:
save_path = '{}/epoch_{}.pth'.format(config.workspace, epoch)
latest_path = '{}/latest.pth'.format(config.workspace)
save_checkpoint(save_path, model, optimizer, epoch, logger)
save_checkpoint(latest_path, model, optimizer, epoch, logger)