print('----------------------------------')
unparsed = vars(ops) # parse_args()方法的返回值为namespace，用vars()内建函数化为字典
for key in unparsed.keys():
print('{} : {}'.format(key,unparsed[key]))
#---------------------------------------------------------------------------
os.environ['CUDA_VISIBLE_DEVICES'] = ops.GPUS
test_path = ops.test_path # 测试图片文件夹路径
#---------------------------------------------------------------- 构建模型
print('use model : %s'%(ops.model))
if ops.model == 'resnet_50':
model_ = resnet50(num_classes = ops.num_classes)
else:
model_ = MY_Net(num_classes = ops.num_classes)
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
model_ = model_.to(device)
model_.eval() # 设置为前向推断模式
# print(model_)# 打印模型结构
# 加载测试模型
if os.access(ops.test_model,os.F_OK):# checkpoint
chkpt = torch.load(ops.test_model, map_location=device)
model_.load_state_dict(chkpt)
def trainer(ops, f_log):
try:
os.environ['CUDA_VISIBLE_DEVICES'] = ops.GPUS
if ops.log_flag:
sys.stdout = f_log
set_seed(ops.seed)
#---------------------------------------------------------------- 构建模型
print('use model : %s' % (ops.model))
if ops.model == 'resnet_50':
model_ = resnet50(num_classes=ops.num_classes,
dropout_factor=ops.dropout)
else:
model_ = MY_Net(num_classes=ops.num_classes)
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
model_ = model_.to(device)
# print(model_)# 打印模型结构
# Dataset
dataset = LoadImagesAndLabels(ops=ops,
img_size=ops.img_size,
flag_agu=ops.flag_agu,
fix_res=ops.fix_res,
vis=False)
print('len train datasets : %s' % (dataset.__len__()))
# Dataloader
dataloader = DataLoader(dataset,
batch_size=ops.batch_size,
num_workers=ops.num_workers,
shuffle=True,
pin_memory=False,
drop_last=True)
# 优化器设计
optimizer_SGD = optim.SGD(model_.parameters(),
lr=ops.init_lr,
momentum=ops.momentum,
weight_decay=ops.weight_decay) # 优化器初始化
optimizer = optimizer_SGD
# 加载 finetune 模型
if os.access(ops.fintune_model, os.F_OK): # checkpoint
chkpt = torch.load(ops.fintune_model, map_location=device)
model_.load_state_dict(chkpt)
print('load fintune model : {}'.format(ops.fintune_model))
print('/**********************************************/')
# 损失函数
if ops.loss_define != 'wing_loss':
criterion = nn.MSELoss(reduce=True, reduction='mean')
step = 0
idx = 0
# 变量初始化
best_loss = np.inf
loss_mean = 0. # 损失均值
loss_idx = 0. # 损失计算计数器
flag_change_lr_cnt = 0 # 学习率更新计数器
init_lr = ops.init_lr # 学习率
epochs_loss_dict = {}
for epoch in range(0, ops.epochs):
if ops.log_flag:
sys.stdout = f_log
print('\nepoch %d ------>>>' % epoch)
model_.train()
# 学习率更新策略
if loss_mean != 0.:
if best_loss > (loss_mean / loss_idx):
flag_change_lr_cnt = 0
best_loss = (loss_mean / loss_idx)
else:
flag_change_lr_cnt += 1
if flag_change_lr_cnt > 20:
init_lr = init_lr * ops.lr_decay
set_learning_rate(optimizer, init_lr)
flag_change_lr_cnt = 0
loss_mean = 0. # 损失均值
loss_idx = 0. # 损失计算计数器
for i, (imgs_, pts_) in enumerate(dataloader):
# print('imgs_, pts_',imgs_.size(), pts_.size())
if use_cuda:
imgs_ = imgs_.cuda(
) # pytorch 的 数据输入格式 ： (batch, channel, height, width)
pts_ = pts_.cuda()
output = model_(imgs_.float())
if ops.loss_define == 'wing_loss':
loss = got_total_wing_loss(output, pts_.float())
else:
loss = criterion(output, pts_.float())
loss_mean += loss.item()
loss_idx += 1.
if i % 10 == 0:
loc_time = time.strftime("%Y-%m-%d %H:%M:%S",
time.localtime())
print(' %s - %s - epoch [%s/%s] (%s/%s):'%(loc_time,ops.model,epoch,ops.epochs,i,int(dataset.__len__()/ops.batch_size)),\
'Mean Loss : %.6f - Loss: %.6f'%(loss_mean/loss_idx,loss.item()),\
' lr : %.5f'%init_lr,' bs :',ops.batch_size,\
' img_size: %s x %s'%(ops.img_size[0],ops.img_size[1]),' best_loss: %.4f'%best_loss)
# 计算梯度
loss.backward()
# 优化器对模型参数更新
optimizer.step()
# 优化器梯度清零
optimizer.zero_grad()
step += 1
torch.save(model_.state_dict(),
ops.model_exp + 'model_epoch-{}.pth'.format(epoch))
except Exception as e:
print('Exception : ', e) # 打印异常
print('Exception file : ',
e.__traceback__.tb_frame.f_globals['__file__']) # 发生异常所在的文件
print('Exception line : ', e.__traceback__.tb_lineno) # 发生异常所在的行数
def trainer(ops, f_log):
try:
os.environ['CUDA_VISIBLE_DEVICES'] = ops.GPUS
if ops.log_flag:
sys.stdout = f_log
set_seed(ops.seed)
train_path = ops.train_path
num_classes = len(os.listdir(ops.train_path)) # 模型类别个数
print('num_classes : ', num_classes)
#---------------------------------------------------------------- 构建模型
print('use model : %s' % (ops.model))
if ops.model == 'resnet_50':
model_ = resnet50(num_classes=num_classes)
else:
model_ = MY_Net(num_classes=num_classes)
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
model_ = model_.to(device)
# print(model_)# 打印模型结构
# Dataset
dataset = LoadImagesAndLabels(path=ops.train_path,
img_size=ops.img_size)
print('len train datasets : %s' % (dataset.__len__()))
# Dataloader
dataloader = DataLoader(dataset,
batch_size=ops.batch_size,
num_workers=ops.num_workers,
shuffle=True,
pin_memory=False,
drop_last=True)
# 优化器设计
# optimizer_Adam = torch.optim.Adam(model_.parameters(), lr=init_lr, betas=(0.9, 0.99),weight_decay=1e-6)
optimizer_SGD = optim.SGD(model_.parameters(),
lr=ops.init_lr,
momentum=0.9,
weight_decay=ops.weight_decay) # 优化器初始化
optimizer = optimizer_SGD
# 加载 finetune 模型
if os.access(ops.fintune_model, os.F_OK): # checkpoint
chkpt = torch.load(ops.fintune_model, map_location=device)
model_.load_state_dict(chkpt)
print('load fintune model : {}'.format(ops.fintune_model))
print('/**********************************************/')
# 损失函数
criterion = nn.CrossEntropyLoss(
) #CrossEntropyLoss() 是 softmax 和 负对数损失的结合
# 变量初始化
best_loss = np.inf
loss_mean = 0. # 损失均值
loss_idx = 0. # 损失计算计数器
flag_change_lr_cnt = 0 # 学习率更新计数器
init_lr = ops.init_lr # 学习率
epochs_loss_dict = {}
for epoch in range(0, ops.epochs):
if ops.log_flag:
sys.stdout = f_log
print('\nepoch %d ------>>>' % epoch)
model_.train()
# 学习率更新策略
if loss_mean != 0.:
if best_loss > (loss_mean / loss_idx):
flag_change_lr_cnt = 0
best_loss = (loss_mean / loss_idx)
else:
flag_change_lr_cnt += 1
if flag_change_lr_cnt > 5:
init_lr = init_lr * ops.lr_decay
set_learning_rate(optimizer, init_lr)
flag_change_lr_cnt = 0
loss_mean = 0. # 损失均值
loss_idx = 0. # 损失计算计数器
for i, (imgs_, labels_) in enumerate(dataloader):
if use_cuda:
imgs_ = imgs_.cuda(
) # pytorch 的 数据输入格式 ： (batch, channel, height, width)
labels_ = labels_.cuda()
output_ = model_(imgs_.float())
loss = criterion(output_, labels_)
loss_mean += loss.item()
loss_idx += 1.
if i % 10 == 0:
acc = get_acc(output_, labels_)
loc_time = time.strftime("%Y-%m-%d %H:%M:%S",
time.localtime())
print(' %s - %s - epoch [%s/%s] (%s/%s):'%(loc_time,ops.model,epoch,ops.epochs,i,int(dataset.__len__()/ops.batch_size)),\
'loss : %.6f - %.6f'%(loss_mean/loss_idx,loss.item()),\
' acc : %.4f'%acc,' lr : %.5f'%init_lr,' bs :',ops.batch_size,\
' img_size: %s x %s'%(ops.img_size[0],ops.img_size[1]),' best_loss: %.4f'%best_loss)
# 计算梯度
loss.backward()
# 优化器对模型参数更新
optimizer.step()
# 优化器梯度清零
optimizer.zero_grad()
# 一个 epoch 保存连词最新的 模型
if i % (int(dataset.__len__() / ops.batch_size / 2 -
1)) == 0 and i > 0:
torch.save(model_.state_dict(),
ops.model_exp + 'latest.pth')
# 每一个 epoch 进行模型保存
torch.save(model_.state_dict(),
ops.model_exp + 'model_epoch-{}.pth'.format(epoch))
except Exception as e:
print('Exception : ', e) # 打印异常
print('Exception file : ',
e.__traceback__.tb_frame.f_globals['__file__']) # 发生异常所在的文件
print('Exception line : ', e.__traceback__.tb_lineno) # 发生异常所在的行数
print('device:', device)
p_model_path = './ckpt/P-Net_latest.pth'
r_model_path = './ckpt/R-Net_latest.pth'
o_model_path = './ckpt/O-Net_latest.pth'
m_PNet, m_RNet, m_ONet = create_XNet(device, p_model_path, r_model_path,
o_model_path)
mtcnn_detector = MtcnnDetector(pnet=m_PNet,
rnet=m_RNet,
onet=m_ONet,
min_face_size=60,
threshold=[0.45, 0.5, 0.75])
#
landmarks_model = resnet50(num_classes=196)
path_landmarks = './ckpt/landmarks.pth'
if os.access(path_landmarks, os.F_OK): # checkpoint
ckpt = torch.load(path_landmarks, map_location=device)
landmarks_model.load_state_dict(ckpt)
print('load landmarks model : {}'.format(path_landmarks))
else:
print('error load landmarks model ')
landmarks_model = landmarks_model.to(device)
landmarks_model.eval() # 设置为前向推断模式
#---------------------------------------------------------------------------
print('init MtcnnDetector')
got_fps = 0.
t_tk_0 = time.time()
show_fps_cnt = 0.