diff --git a/train.py b/train.py index b4f9de7f0efd0d1a61c76396bfecbc50d9ed7032..5277f3b939f1e9b07314e4fdb2b08ce19f09939d 100644 --- a/train.py +++ b/train.py @@ -101,8 +101,6 @@ def trainer(ops,f_log): # 损失函数 if ops.loss_define == 'mse_loss': criterion = nn.MSELoss(reduce=True, reduction='mean') - elif ops.loss_define == 'adaptive_wing_loss': - criterion = AdaptiveWingLoss() step = 0 idx = 0 @@ -195,8 +193,8 @@ if __name__ == "__main__": help = 'imageNet_Pretrain') # 初始化学习率 parser.add_argument('--fintune_model', type=str, default = 'None', help = 'fintune_model') # fintune model - parser.add_argument('--loss_define', type=str, default = 'adaptive_wing_loss', - help = 'define_loss : wing_loss, mse_loss ,adaptive_wing_loss') # 损失函数定义 + parser.add_argument('--loss_define', type=str, default = 'wing_loss', + help = 'define_loss : wing_loss, mse_loss ') # 损失函数定义 parser.add_argument('--init_lr', type=float, default = 1e-3, help = 'init learning Rate') # 初始化学习率 parser.add_argument('--lr_decay', type=float, default = 0.1,