提交 56ca4526 编写于 作者: Eric.Lee2021's avatar Eric.Lee2021 🚴🏻

Update train.py

上级 da7ef089
...@@ -101,8 +101,6 @@ def trainer(ops,f_log): ...@@ -101,8 +101,6 @@ def trainer(ops,f_log):
# 损失函数 # 损失函数
if ops.loss_define == 'mse_loss': if ops.loss_define == 'mse_loss':
criterion = nn.MSELoss(reduce=True, reduction='mean') criterion = nn.MSELoss(reduce=True, reduction='mean')
elif ops.loss_define == 'adaptive_wing_loss':
criterion = AdaptiveWingLoss()
step = 0 step = 0
idx = 0 idx = 0
...@@ -195,8 +193,8 @@ if __name__ == "__main__": ...@@ -195,8 +193,8 @@ if __name__ == "__main__":
help = 'imageNet_Pretrain') # 初始化学习率 help = 'imageNet_Pretrain') # 初始化学习率
parser.add_argument('--fintune_model', type=str, default = 'None', parser.add_argument('--fintune_model', type=str, default = 'None',
help = 'fintune_model') # fintune model help = 'fintune_model') # fintune model
parser.add_argument('--loss_define', type=str, default = 'adaptive_wing_loss', parser.add_argument('--loss_define', type=str, default = 'wing_loss',
help = 'define_loss : wing_loss, mse_loss ,adaptive_wing_loss') # 损失函数定义 help = 'define_loss : wing_loss, mse_loss ') # 损失函数定义
parser.add_argument('--init_lr', type=float, default = 1e-3, parser.add_argument('--init_lr', type=float, default = 1e-3,
help = 'init learning Rate') # 初始化学习率 help = 'init learning Rate') # 初始化学习率
parser.add_argument('--lr_decay', type=float, default = 0.1, parser.add_argument('--lr_decay', type=float, default = 0.1,
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册