def __init__(self, box_codec, anchors_cnt=6, classes_cnt=21, classification_weight=1., regression_weight=1.):
super(RetinaNetLoss, self).__init__()
self.anchors_cnt = anchors_cnt
self.classes_cnt = classes_cnt
self.box_codec = box_codec
self.classifcation_w = classification_weight
self.regression_w = regression_weight
self.regression_criterion = SmoothL1Loss()
self.classification_criterion = FocalLoss()
def __init__(self):
# Action definition
self.USE_CUDA = torch.cuda.is_available()
self.LOAD_SAVED_MOD = True
self.SAVE_TEMP_MODEL = True
self.SAVE_BEST_MODEL = True
self.BEST_MODEL_BY_LOSS = False
self.PRINT_BAD_CASE = True
self.RUNNING_ON_JUPYTER = False
self.START_VOTE_PREDICT = False
self.START_PREDICT = False
self.TRAIN_ALL = False
self.TEST_ALL = False
self.TO_MULTI = False
self.ADD_SUMMARY = False
self.SAVE_PER_EPOCH = 1
# Tensor shape definition
self.BATCH_SIZE = 32
self.VAL_BATCH_SIZE = 64
self.TENSOR_SHAPE = (3, 224, 224)
# Program information
self.DATALOADER_TYPE = "ImageFolder"
self.OPTIMIZER = "Adam"
self.SGD_MOMENTUM = 0.9
self.OPT_WEIGHT_DECAY = 0
self.TRAIN_DATA_RATIO = 0.7
self.NUM_EPOCHS = 500
self.NUM_CLASSES = 250
self.NUM_VAL = 1
self.NUM_TRAIN = 1
self.TOP_NUM = 1
self.NUM_WORKERS = 0
self.CRITERION = FocalLoss(
self.NUM_CLASSES) #torch.nn.CrossEntropyLoss()
# Hyper parameters
self.LEARNING_RATE = 0.0001
self.TOP_VOTER = 6
# Name and path definition
self.NET_SAVE_PATH = "./source/trained_net/"
self.TRAIN_PATH = "../cards_for_train"
self.VAL_PATH = "../cards_250_7/cards_for_val"
self.CLASSES_PATH = "./reference/classes_name.pkl"
self.MODEL_NAME = "MobileNetV2"
self.PROCESS_ID = "Test05_250_NewDataset" #_FocalLoss"
if self.TRAIN_ALL:
self.PROCESS_ID += '_TRAIN_ALL'
self.SUMMARY_PATH = "./source/summary/" + self.MODEL_NAME + '_' + self.PROCESS_ID
def test_focal_loss():
alpha, gamma = 0.25, 2
fl = FocalLoss(alpha=alpha, gamma=gamma)
logits = torch.from_numpy(np.asarray([[0.05, 0.3, -2], [0.5, -0.01, 0.03]], dtype=np.float32))
target = torch.from_numpy(np.asarray([[0, 1, 0], [0, 0, 1]], dtype=np.float32))
computed_loss = fl(pred_logits=logits, target=target)
loss = torch.sum(computed_loss).item()
target_loss = 0.
for batch, batch_tgt in zip(logits, target):
probs = torch.sigmoid(batch)
for p, t in zip(probs, batch_tgt):
pt = p if t else 1. - p
alpha_t = alpha if t else 1. - alpha
target_loss += -1. * alpha_t * ((1. - pt) ** gamma) * torch.log(pt)
target_loss = target_loss.item()
assert loss == pytest.approx(target_loss, 1e-3)
def train(args, model, optimizer, data_loader):
model.train()
min_loss = 1e6
no_implove_cnt = 0
criteria = FocalLoss(args.n_class, ignore_label=-1)
for epoch in range(args.epochs):
for i, (l_data, l_target) in enumerate(data_loader):
l_data = l_data.to(device)
l_target = l_target.to(device)
model.zero_grad()
optimizer.zero_grad()
output = model(l_data)
targets = l_target.view(-1)
loss = criteria(output, targets)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
optimizer.step()
print('[{}/{}][{}/{}] Loss: {:.4f}'.format(epoch, args.epochs, i,
len(data_loader),
loss.item()))
min_loss = loss.item() if min_loss > loss.item() else min_loss
no_implove_cnt = no_implove_cnt + 1 if loss.item(
) > min_loss else 0
if no_implove_cnt == 250:
exp_lr_scheduler(optimizer, epoch, init_lr=args.lr)
no_implove_cnt = 0
# do checkpointing
if epoch % 5 == 0:
torch.save(
model.state_dict(),
'{}/{}_model_ckpt.pth'.format(args.out_dir, args.target))
torch.save(model.state_dict(),
'{}/{}_model_ckpt.pth'.format(args.out_dir, args.target))
def train(Config,
model,
epoch_num,
start_epoch,
optimizer,
exp_lr_scheduler,
data_loader,
save_dir,
data_size=448,
savepoint=500,
checkpoint=1000):
# savepoint: save without evalution
# checkpoint: save with evaluation
step = 0
eval_train_flag = False
rec_loss = []
checkpoint_list = []
train_batch_size = data_loader['train'].batch_size
train_epoch_step = data_loader['train'].__len__()
train_loss_recorder = LossRecord(train_batch_size)
if savepoint > train_epoch_step:
savepoint = 1 * train_epoch_step
checkpoint = savepoint
date_suffix = dt()
log_file = open(
os.path.join(
Config.log_folder,
'formal_log_r50_dcl_%s_%s.log' % (str(data_size), date_suffix)),
'a')
add_loss = nn.L1Loss()
get_ce_loss = nn.CrossEntropyLoss()
get_focal_loss = FocalLoss()
get_angle_loss = AngleLoss()
for epoch in range(start_epoch, epoch_num - 1):
exp_lr_scheduler.step(epoch)
model.train(True)
save_grad = []
for batch_cnt, data in enumerate(data_loader['train']):
step += 1
loss = 0
model.train(True)
if Config.use_backbone:
inputs, labels, img_names = data
inputs = Variable(inputs.cuda())
labels = Variable(torch.from_numpy(np.array(labels)).cuda())
if Config.use_dcl:
inputs, labels, labels_swap, swap_law, img_names = data
inputs = Variable(inputs.cuda())
labels = Variable(torch.from_numpy(np.array(labels)).cuda())
labels_swap = Variable(
torch.from_numpy(np.array(labels_swap)).cuda())
swap_law = Variable(
torch.from_numpy(np.array(swap_law)).float().cuda())
optimizer.zero_grad()
if inputs.size(0) < 2 * train_batch_size:
outputs = model(inputs, inputs[0:-1:2])
else:
outputs = model(inputs, None)
if Config.use_focal_loss:
ce_loss = get_focal_loss(outputs[0], labels)
else:
ce_loss = get_ce_loss(outputs[0], labels)
if Config.use_Asoftmax:
fetch_batch = labels.size(0)
if batch_cnt % (train_epoch_step // 5) == 0:
angle_loss = get_angle_loss(outputs[3],
labels[0:fetch_batch:2],
decay=0.9)
else:
angle_loss = get_angle_loss(outputs[3],
labels[0:fetch_batch:2])
loss += angle_loss
loss += ce_loss
alpha_ = 1
beta_ = 1
gamma_ = 0.01 if Config.dataset == 'STCAR' or Config.dataset == 'AIR' else 1
if Config.use_dcl:
swap_loss = get_ce_loss(outputs[1], labels_swap) * beta_
loss += swap_loss
law_loss = add_loss(outputs[2], swap_law) * gamma_
loss += law_loss
loss.backward()
torch.cuda.synchronize()
optimizer.step()
torch.cuda.synchronize()
if Config.use_dcl:
print(
'step: {:-8d} / {:d} loss=ce_loss+swap_loss+law_loss: {:6.4f} = {:6.4f} + {:6.4f} + {:6.4f} '
.format(step, train_epoch_step,
loss.detach().item(),
ce_loss.detach().item(),
swap_loss.detach().item(),
law_loss.detach().item()),
flush=True)
if Config.use_backbone:
print(
'step: {:-8d} / {:d} loss=ce_loss+swap_loss+law_loss: {:6.4f} = {:6.4f} '
.format(step, train_epoch_step,
loss.detach().item(),
ce_loss.detach().item()),
flush=True)
rec_loss.append(loss.detach().item())
train_loss_recorder.update(loss.detach().item())
# evaluation & save
if step % checkpoint == 0:
rec_loss = []
print(32 * '-', flush=True)
print(
'step: {:d} / {:d} global_step: {:8.2f} train_epoch: {:04d} rec_train_loss: {:6.4f}'
.format(step, train_epoch_step,
1.0 * step / train_epoch_step, epoch,
train_loss_recorder.get_val()),
flush=True)
print('current lr:%s' % exp_lr_scheduler.get_lr(), flush=True)
if eval_train_flag:
trainval_acc1, trainval_acc2, trainval_acc3 = eval_turn(
model, data_loader['trainval'], 'trainval', epoch,
log_file)
if abs(trainval_acc1 - trainval_acc3) < 0.01:
eval_train_flag = False
val_acc1, val_acc2, val_acc3 = eval_turn(
model, data_loader['val'], 'val', epoch, log_file)
save_path = os.path.join(
save_dir, 'weights_%d_%d_%.4f_%.4f.pth' %
(epoch, batch_cnt, val_acc1, val_acc3))
torch.cuda.synchronize()
torch.save(model.state_dict(), save_path)
print('saved model to %s' % (save_path), flush=True)
torch.cuda.empty_cache()
# save only
elif step % savepoint == 0:
train_loss_recorder.update(rec_loss)
rec_loss = []
save_path = os.path.join(
save_dir, 'savepoint_weights-%d-%s.pth' % (step, dt()))
checkpoint_list.append(save_path)
if len(checkpoint_list) == 6:
os.remove(checkpoint_list[0])
del checkpoint_list[0]
torch.save(model.state_dict(), save_path)
torch.cuda.empty_cache()
log_file.close()
shuffle=True,
num_workers=opt.num_workers)
val_dataset = Dataset(opt.train_root, opt.val_list, phase='val', input_shape=opt.input_shape)
valloader = data.DataLoader(val_dataset,
batch_size=opt.train_batch_size,
shuffle=False,
num_workers=opt.num_workers
)
identity_list = get_lfw_list(opt.lfw_test_list)
img_paths = [os.path.join(opt.lfw_root, each) for each in identity_list]
print('{} train iters per epoch:'.format(len(trainloader)))
if opt.loss == 'focal_loss':
criterion = FocalLoss(gamma=2)
else:
criterion = torch.nn.CrossEntropyLoss()
if opt.backbone == 'resnet_face18':
model = resnet_face18(opt.input_shape, use_se=opt.use_se)
elif opt.backbone == 'resnet34':
model = resnet34(input_shape=opt.input_shape)
elif opt.backbone == 'resnet50':
model = resnet50(input_shape=opt.input_shape)
elif opt.backbone == 'resnet_face50':
model = resnet_face50(opt.input_shape, use_se=opt.use_se)
elif opt.backbone == 'resnet_face100':
model = resnet_face100(opt.input_shape, use_se=opt.use_se)
elif opt.backbone == 'resnet101':
model = resnet101(input_shape=opt.input_shape)
def forward(self, predictions, targets, anchors, anchor_t, gr):
device = targets.device
lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device)
#构建目标
tcls, tbox, indices, anchors = build_targets(predictions, targets, anchors, anchor_t)
hyp = config.hyp
# Define criteria
BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.Tensor([hyp['cls_pw']])).to(device)
BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.Tensor([hyp['obj_pw']])).to(device)
# Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
cp, cn = smooth_BCE(eps=0.0)
# Focal loss
g = hyp['fl_gamma'] # focal loss gamma
if g > 0:
BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
# Losses
nt = 0 # number of targets
np = len(predictions) # number of outputs
balance = [4.0, 1.0, 0.4] if np == 3 else [4.0, 1.0, 0.4, 0.1] # P3-5 or P3-6
for i, pi in enumerate(predictions): # layer index, layer predictions
b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
tobj = torch.zeros_like(pi[..., 0], device=device) # target obj
n = b.shape[0] # number of targets
if n:
nt += n # cumulative targets
ps = pi[b, a, gj, gi] # prediction subset corresponding to targets
# Regression
pxy = ps[:, :2].sigmoid() * 2. - 0.5
pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
pbox = torch.cat((pxy, pwh), 1).to(device) # predicted box
giou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True) # giou(prediction, target)
lbox += (1.0 - giou).mean() # giou loss
# Objectness
tobj[b, a, gj, gi] = (1.0 - gr) + gr * giou.detach().clamp(0).type(tobj.dtype) # giou ratio
# Classification
if self.num_class > 1: # cls loss (only if multiple classes)
t = torch.full_like(ps[:, 5:], cn, device=device) # targets
t[range(n), tcls[i]] = cp
lcls += BCEcls(ps[:, 5:], t) # BCE
# Append targets to text file
# with open('targets.txt', 'a') as file:
# [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]
lobj += BCEobj(pi[..., 4], tobj) * balance[i] # obj loss
s = 3 / np # output count scaling
lbox *= hyp['giou'] * s
lobj *= hyp['obj'] * s * (1.4 if np == 4 else 1.)
lcls *= hyp['cls'] * s
bs = tobj.shape[0] # batch size
loss = lbox + lobj + lcls
return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()
def main():
os.makedirs(args.checkpoints, exist_ok=True)
torch.manual_seed(7)
np.random.seed(7)
best_acc = 0
if args.resume_checkpoint is None:
if args.backbone == 'r18':
model = resnet18(args)
elif args.backbone == 'r34':
model = resnet34(args)
elif args.backbone == 'r50':
model = resnet50(args)
elif args.backbone == 'resnet101':
model = resnet101(args)
elif args.backbone == 'r152':
model = resnet152(args)
elif args.backbone == 'mobile':
model = MobileNet(1.0)
else:
model = resnet_face18(args.use_se)
model = nn.DataParallel(model)
num_person = calc_num_person(args.train_path)
print('======================================')
print('To be trained people number is: {}.'.format(num_person))
print('======================================')
metric_fc = ArcMarginModel(args, num_person)
metric_fc = nn.DataParallel(metric_fc)
if args.optimizer == 'sgd':
optimizer = InsightFaceOptimizer(
torch.optim.SGD([{'params': model.parameters()}, {'params': metric_fc.parameters()}],
lr=args.lr, momentum=args.mom, weight_decay=args.weight_decay))
else:
optimizer = torch.optim.Adam([{'params': model.parameters()}, {'params': metric_fc.parameters()}],
lr=args.lr, weight_decay=args.weight_decay)
else:
checkpoint = torch.load(args.resume_checkpoint)
model = checkpoint['model']
metric_fc = checkpoint['metric_fc']
optimizer = checkpoint['optimizer']
model = model.to(device)
metric_fc = metric_fc.to(device)
# Loss function
if args.focal_loss:
criterion = FocalLoss(gamma=args.gamma).to(device)
else:
criterion = nn.CrossEntropyLoss().to(device)
train_dataset = Dataset(root=args.train_path, phase='train', input_shape=(3, 112, 112))
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=8)
for epoch in range(args.epochs):
loss, acc = train_one_epoch(train_loader, model, metric_fc, criterion, optimizer, epoch)
if acc > best_acc:
status = {'epoch': epoch, 'acc': acc, 'model': model, 'metric_fc': metric_fc, 'optimizer': optimizer}
torch.save(status, os.path.join(args.checkpoints, str(epoch) + '_' + '%.4f' % acc + '.pth'))
best_acc = acc
def run():
opt = Config()
if opt.display:
visualizer = Visualizer()
# device = torch.device("cuda")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_dataset = FaceDataset(opt.train_root,
opt.train_list,
phase='train',
input_shape=opt.input_shape)
trainloader = torch.utils.data.DataLoader(train_dataset,
batch_size=opt.train_batch_size,
shuffle=True,
num_workers=opt.num_workers)
print('{} train iters per epoch:'.format(len(trainloader)))
# Focal Loss, 解决类别不均衡问题,减少易分类样本的权重，使得模型在训练时更专注于难分类的样本
# https://blog.csdn.net/u014380165/article/details/77019084
#
#定义损失函数
if opt.loss == 'focal_loss':
criterion = FocalLoss(gamma=2) #
else:
criterion = torch.nn.CrossEntropyLoss()
#定义模型
if opt.backbone == 'resnet18':
model = resnet_face18(use_se=opt.use_se)
elif opt.backbone == 'resnet34':
model = resnet34()
elif opt.backbone == 'resnet50':
model = resnet50()
#全连接层？
if opt.metric == 'add_margin':
metric_fc = AddMarginProduct(512, opt.num_classes, s=30, m=0.35)
elif opt.metric == 'arc_margin':
metric_fc = ArcMarginProduct(512,
opt.num_classes,
s=30,
m=0.5,
easy_margin=opt.easy_margin)
elif opt.metric == 'sphere':
metric_fc = SphereProduct(512, opt.num_classes, m=4)
else:
metric_fc = nn.Linear(512, opt.num_classes)
# view_model(model, opt.input_shape)
print(model)
model.to(device)
model = DataParallel(model)
metric_fc.to(device)
metric_fc = DataParallel(metric_fc)
#定义优化算法
if opt.optimizer == 'sgd':
optimizer = torch.optim.SGD([{
'params': model.parameters()
}, {
'params': metric_fc.parameters()
}],
lr=opt.lr,
weight_decay=opt.weight_decay)
else:
optimizer = torch.optim.Adam([{
'params': model.parameters()
}, {
'params': metric_fc.parameters()
}],
lr=opt.lr,
weight_decay=opt.weight_decay)
# https://www.programcreek.com/python/example/98143/torch.optim.lr_scheduler.StepLR
# ? 每过{lr_step}个epoch训练，学习率就乘gamma
scheduler = StepLR(optimizer, step_size=opt.lr_step, gamma=0.1)
start = time.time()
for i in range(opt.max_epoch):
scheduler.step()
model.train() # train模式，eval模式
for ii, data in enumerate(trainloader):
data_input, label = data
data_input = data_input.to(device)
label = label.to(device).long()
feature = model(data_input)
output = metric_fc(feature,
label) # 全连接层？ 将原本用于输出分类的层，改成输出512维向量？似乎不是？
loss = criterion(output, label) # criterion:做出判断的依据
optimizer.zero_grad()
loss.backward()
optimizer.step()
iters = i * len(trainloader) + ii
if iters % opt.print_freq == 0:
output = output.data.cpu().numpy()
output = np.argmax(output,
axis=1) #最大值所在的索引？ index <-> one-hot相互转换
label = label.data.cpu().numpy()
# print(output)
# print(label)
acc = np.mean((output == label).astype(int))
speed = opt.print_freq / (time.time() - start)
time_str = time.asctime(time.localtime(time.time()))
print('{} train epoch {} iter {} {} iters/s loss {} acc {}'.
format(time_str, i, ii, speed, loss.item(), acc))
if opt.display:
visualizer.display_current_results(iters,
loss.item(),
name='train_loss')
visualizer.display_current_results(iters,
acc,
name='train_acc')
start = time.time()
if i % opt.save_interval == 0 or i == opt.max_epoch:
save_model(model, opt.checkpoints_path, opt.backbone, i)
# train结束，模型设置为eval模式
model.eval()
#测试？
identity_list = get_lfw_list(opt.lfw_test_list)
img_paths = [
os.path.join(opt.lfw_root, each) for each in identity_list
]
acc = lfw_test(model, img_paths, identity_list, opt.lfw_test_list,
opt.test_batch_size)
if opt.display:
visualizer.display_current_results(iters, acc, name='test_acc')
def train(Config,
model,
epoch_num,
start_epoch,
optimizer,
exp_lr_scheduler,
data_loader,
save_dir,
data_size=448,
savepoint=500,
checkpoint=1000):
# savepoint: save without evalution
# checkpoint: save with evaluation
bmy_weight = 1.0 # 1.5 # 决定品牌分支在学习中权重
step = 0
eval_train_flag = False
rec_loss = []
checkpoint_list = []
steps = np.array([], dtype=np.int)
train_accs = np.array([], dtype=np.float32)
test_accs = np.array([], dtype=np.float32)
ce_losses = np.array([], dtype=np.float32)
ce_loss_mu = -1
ce_loss_std = 0.0
train_batch_size = data_loader['train'].batch_size
train_epoch_step = data_loader['train'].__len__()
train_loss_recorder = LossRecord(train_batch_size)
if savepoint > train_epoch_step:
savepoint = 1 * train_epoch_step
checkpoint = savepoint
date_suffix = dt()
log_file = open(
os.path.join(
Config.log_folder,
'formal_log_r50_dcl_%s_%s.log' % (str(data_size), date_suffix)),
'a')
add_loss = nn.L1Loss()
get_ce_loss = nn.CrossEntropyLoss()
get_focal_loss = FocalLoss()
get_angle_loss = AngleLoss()
for epoch in range(start_epoch, epoch_num - 1):
model.train(True)
save_grad = []
for batch_cnt, data in enumerate(data_loader['train']):
step += 1
loss = 0
model.train(True)
if Config.use_backbone:
inputs, brand_labels, img_names, bmy_labels = data
inputs = Variable(inputs.cuda())
brand_labels = Variable(
torch.from_numpy(np.array(brand_labels)).cuda())
bmy_labels = Variable(
torch.from_numpy(np.array(bmy_labels)).cuda())
if Config.use_dcl:
inputs, brand_labels, brand_labels_swap, swap_law, img_names, bmy_labels = data
org_brand_labels = brand_labels
inputs = Variable(inputs.cuda())
brand_labels = Variable(
torch.from_numpy(np.array(brand_labels)).cuda())
bmy_labels = Variable(
torch.from_numpy(np.array(bmy_labels)).cuda())
brand_labels_swap = Variable(
torch.from_numpy(np.array(brand_labels_swap)).cuda())
swap_law = Variable(
torch.from_numpy(np.array(swap_law)).float().cuda())
optimizer.zero_grad()
if inputs.size(0) < 2 * train_batch_size:
outputs = model(inputs, inputs[0:-1:2])
else:
outputs = model(inputs, None)
if Config.use_focal_loss:
ce_loss_brand = get_focal_loss(outputs[0], brand_labels)
ce_loss_bmy = get_focal_loss(outputs[-1], bmy_labels)
else:
ce_loss_brand = get_ce_loss(outputs[0], brand_labels)
ce_loss_bmy = get_ce_loss(outputs[-1], bmy_labels)
ce_loss = ce_loss_brand + bmy_weight * ce_loss_bmy
if Config.use_Asoftmax:
fetch_batch = brand_labels.size(0)
if batch_cnt % (train_epoch_step // 5) == 0:
angle_loss = get_angle_loss(outputs[3],
brand_labels[0:fetch_batch:2],
decay=0.9)
else:
angle_loss = get_angle_loss(outputs[3],
brand_labels[0:fetch_batch:2])
loss += angle_loss
loss += ce_loss
ce_loss_val = ce_loss.detach().item()
ce_losses = np.append(ce_losses, ce_loss_val)
alpha_ = 1
beta_ = 1
gamma_ = 0.01 if Config.dataset == 'STCAR' or Config.dataset == 'AIR' else 1
if Config.use_dcl:
swap_loss = get_ce_loss(outputs[1], brand_labels_swap) * beta_
loss += swap_loss
law_loss = add_loss(outputs[2], swap_law) * gamma_
loss += law_loss
loss.backward()
torch.cuda.synchronize()
optimizer.step()
exp_lr_scheduler.step(epoch)
torch.cuda.synchronize()
if Config.use_dcl:
if ce_loss_mu > 0 and ce_loss_val > ce_loss_mu + 3.0 * ce_loss_std:
# 记录下这个批次，可能是该批次有标注错误情况
print('记录可疑批次信息: loss={0}; threshold={1};'.format(
ce_loss_val, ce_loss_mu + 2.0 * ce_loss_std))
with open(
'./logs/abnormal_samples_{0}_{1}_{2}.txt'.format(
epoch, step, ce_loss_val), 'a+') as fd:
error_batch_len = len(img_names)
for i in range(error_batch_len):
fd.write('{0} <=> {1};\r\n'.format(
org_brand_labels[i * 2], img_names[i]))
print('epoch{}: step: {:-8d} / {:d} loss=ce_loss+'
'swap_loss+law_loss: {:6.4f} = {:6.4f} '
'+ {:6.4f} + {:6.4f} brand_loss: {:6.4f}'.format(
epoch, step % train_epoch_step, train_epoch_step,
loss.detach().item(), ce_loss_val,
swap_loss.detach().item(),
law_loss.detach().item(),
ce_loss_brand.detach().item()),
flush=True)
if Config.use_backbone:
print('epoch{}: step: {:-8d} / {:d} loss=ce_loss+'
'swap_loss+law_loss: {:6.4f} = {:6.4f} '.format(
epoch, step % train_epoch_step, train_epoch_step,
loss.detach().item(),
ce_loss.detach().item()),
flush=True)
rec_loss.append(loss.detach().item())
train_loss_recorder.update(loss.detach().item())
# evaluation & save
if step % checkpoint == 0:
rec_loss = []
print(32 * '-', flush=True)
print(
'step: {:d} / {:d} global_step: {:8.2f} train_epoch: {:04d} rec_train_loss: {:6.4f}'
.format(step, train_epoch_step,
1.0 * step / train_epoch_step, epoch,
train_loss_recorder.get_val()),
flush=True)
print('current lr:%s' % exp_lr_scheduler.get_lr(), flush=True)
'''
if eval_train_flag:
trainval_acc1, trainval_acc2, trainval_acc3 = eval_turn(Config, model, data_loader['trainval'], 'trainval', epoch, log_file)
if abs(trainval_acc1 - trainval_acc3) < 0.01:
eval_train_flag = False
'''
print('##### validate dataset #####')
trainval_acc1, trainval_acc2, trainval_acc3 = eval_turn(
Config, model, data_loader['val'], 'val', epoch, log_file
) #eval_turn(Config, model, data_loader['trainval'], 'trainval', epoch, log_file)
print('##### test dataset #####')
val_acc1, val_acc2, val_acc3 = trainval_acc1, trainval_acc2, \
trainval_acc3 # eval_turn(Config, model, data_loader['val'], 'val', epoch, log_file)
steps = np.append(steps, step)
train_accs = np.append(train_accs, trainval_acc1)
test_accs = np.append(test_accs, val_acc1)
save_path = os.path.join(
save_dir, 'weights_%d_%d_%.4f_%.4f.pth' %
(epoch, batch_cnt, val_acc1, val_acc3))
torch.cuda.synchronize()
torch.save(model.state_dict(),
save_path,
_use_new_zipfile_serialization=False)
print('saved model to %s' % (save_path), flush=True)
torch.cuda.empty_cache()
# 保存精度等信息并初始化
ce_loss_mu = ce_losses.mean()
ce_loss_std = ce_losses.std()
print('Cross entropy loss: mu={0}; std={1}; range:{2}~{3};'.
format(ce_loss_mu, ce_loss_std,
ce_loss_mu - 3.0 * ce_loss_std,
ce_loss_mu + 3.0 * ce_loss_std))
ce_losses = np.array([], dtype=np.float32)
if train_accs.shape[0] > 30:
np.savetxt('./logs/steps1.txt', (steps, ))
np.savetxt('./logs/train_accs1.txt', (train_accs, ))
np.savetxt('./logs/test_accs1.txt', (test_accs, ))
steps = np.array([], dtype=np.int)
train_accs = np.array([], dtype=np.float32)
test_accs = np.array([], dtype=np.float32)
# save only
elif step % savepoint == 0:
train_loss_recorder.update(rec_loss)
rec_loss = []
save_path = os.path.join(
save_dir, 'savepoint_weights-%d-%s.pth' % (step, dt()))
checkpoint_list.append(save_path)
if len(checkpoint_list) == 6:
os.remove(checkpoint_list[0])
del checkpoint_list[0]
torch.save(model.state_dict(),
save_path,
_use_new_zipfile_serialization=False)
torch.cuda.empty_cache()
log_file.close()
args.scheduler = lr_scheduler.ReduceLROnPlateau(args.optimizer,
'min',
patience=10,
factor=0.1)
# load data
train_loader, train_ratio_list, val_loader, val_ratio_list = _get_dataloader(
batch_size=args.batch_size, random_seed=1)
alpha_list = _get_alpha(train_ratio_list)
alpha_list = [1, 1, 20]
args.num_data = {
'train': float(sum(train_ratio_list)),
'val': float(sum(val_ratio_list))
}
if args.loss == 'focal':
args.criterion = FocalLoss(gamma=3, alpha=alpha_list)
elif args.loss == 'nll':
weight = torch.FloatTensor(
alpha_list).cuda() if gpu else torch.FloatTensor(alpha_list)
args.criterion = lambda x, y, weight=weight: nn.NLLLoss(weight=weight)(
F.log_softmax(x), y)
else:
print('undefined loss')
args.writer = SummaryWriter(args.save)
for epoch in range(1, args.epoch):
train(args, net161, epoch, train_loader, gpu=True)
val_loss = val(args, net161, epoch, val_loader, gpu=True)
args.scheduler.step(val_loss)
args.writer.close()
def main():
log_dir = os.path.join('logs', '000')
opt = Config()
if opt.display:
visualizer = Visualizer()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.manual_seed(1)
train_dataset = Dataset(opt.train_root,
opt.train_list,
phase='train',
input_shape=opt.input_shape)
trainloader = data.DataLoader(train_dataset,
batch_size=opt.train_batch_size,
shuffle=True,
num_workers=opt.num_workers)
#identity_list = get_lfw_list(opt.lfw_test_list)
#img_paths = [os.path.join(opt.lfw_root, each) for each in identity_list]
if opt.loss == 'focal_loss':
criterion = FocalLoss(gamma=2)
else:
criterion = torch.nn.CrossEntropyLoss()
if opt.backbone == 'resnet18':
model = resnet_face18(use_se=opt.use_se)
elif opt.backbone == 'resnet34':
#model = resnet34()
model = resnet_face34(use_se=opt.use_se)
elif opt.backbone == 'resnet50':
model = resnet50()
if opt.metric == 'add_margin':
metric_fc = AddMarginProduct(512,
opt.num_classes,
s=30,
m=0.35,
device=device)
elif opt.metric == 'arc_margin':
metric_fc = ArcMarginProduct(512,
opt.num_classes,
s=30,
m=0.5,
easy_margin=opt.easy_margin,
device=device)
elif opt.metric == 'sphere':
metric_fc = SphereProduct(512, opt.num_classes, m=4, device=device)
else:
metric_fc = torch.nn.Linear(512, opt.num_classes)
# view_model(model, opt.input_shape)
#print(model)
model.to(device)
summary(model, input_size=opt.input_shape)
model = DataParallel(model)
metric_fc.to(device)
metric_fc = DataParallel(metric_fc)
print('{} train iters per epoch:'.format(len(trainloader)))
if opt.optimizer == 'sgd':
optimizer = torch.optim.SGD([{
'params': model.parameters()
}, {
'params': metric_fc.parameters()
}],
lr=opt.lr,
weight_decay=opt.weight_decay)
else:
optimizer = torch.optim.Adam([{
'params': model.parameters()
}, {
'params': metric_fc.parameters()
}],
lr=opt.lr,
weight_decay=opt.weight_decay)
scheduler = StepLR(optimizer, step_size=opt.lr_step, gamma=0.1)
#start = time.time()
for epoch in range(opt.max_epoch):
scheduler.step()
print('Epoch %d/%d' % (epoch, opt.max_epoch))
train(opt, model, metric_fc, device, trainloader, criterion, optimizer,
scheduler)
validate(opt, model, device, epoch, log_dir)
def train(Config,
model,
epoch_num,
start_epoch,
optimizer,
exp_lr_scheduler,
data_loader,
save_dir,
sw,
data_size=448,
savepoint=500,
checkpoint=1000):
# savepoint: save without evalution
# checkpoint: save with evaluation
best_prec1 = 0.
step = 0
eval_train_flag = False
rec_loss = []
checkpoint_list = []
train_batch_size = data_loader['train'].batch_size
train_epoch_step = data_loader['train'].__len__()
train_loss_recorder = LossRecord(train_batch_size)
if savepoint > train_epoch_step:
savepoint = 1 * train_epoch_step
checkpoint = savepoint
date_suffix = dt()
# log_file = open(os.path.join(Config.log_folder, 'formal_log_r50_dcl_%s_%s.log'%(str(data_size), date_suffix)), 'a')
add_loss = nn.L1Loss()
get_ce_loss = nn.CrossEntropyLoss()
get_loss1 = Loss_1()
get_focal_loss = FocalLoss()
get_angle_loss = AngleLoss()
for epoch in range(start_epoch, epoch_num - 1):
optimizer.step()
exp_lr_scheduler.step(epoch)
model.train(True)
save_grad = []
for batch_cnt, data in enumerate(data_loader['train']):
step += 1
loss = 0
model.train(True)
if Config.use_backbone:
inputs, labels, img_names = data
inputs = Variable(inputs.cuda())
labels = Variable(torch.from_numpy(np.array(labels)).cuda())
if Config.use_dcl:
if Config.multi:
inputs, labels, labels_swap, swap_law, blabels, clabels, tlabels, img_names = data
else:
inputs, labels, labels_swap, swap_law, img_names = data
inputs = Variable(inputs.cuda())
labels = Variable(torch.from_numpy(np.array(labels)).cuda())
labels_swap = Variable(
torch.from_numpy(np.array(labels_swap)).cuda())
swap_law = Variable(
torch.from_numpy(np.array(swap_law)).float().cuda())
if Config.multi:
blabels = Variable(
torch.from_numpy(np.array(blabels)).cuda())
clabels = Variable(
torch.from_numpy(np.array(clabels)).cuda())
tlabels = Variable(
torch.from_numpy(np.array(tlabels)).cuda())
optimizer.zero_grad()
# 显示输入图片
# sw.add_image('attention_image', inputs[0])
if inputs.size(0) < 2 * train_batch_size:
outputs = model(inputs, inputs[0:-1:2])
else:
outputs = model(inputs, None)
if Config.multi:
if Config.use_loss1:
b_loss, pro_b = get_loss1(outputs[2], blabels)
# 关联品牌标签和车型
t_loss, _ = get_loss1(outputs[4],
tlabels,
brand_prob=pro_b)
s_loss, pro_s = get_loss1(outputs[0],
labels,
brand_prob=pro_b)
c_loss, _ = get_loss1(outputs[3], clabels)
ce_loss = b_loss + t_loss + s_loss + c_loss * 0.2
else:
ce_loss = get_ce_loss(outputs[0], labels) + get_ce_loss(
outputs[0], blabels) + get_ce_loss(
outputs[0], clabels) + get_ce_loss(
outputs[0], tlabels)
else:
if Config.use_focal_loss:
ce_loss = get_focal_loss(outputs[0], labels)
else:
if Config.use_loss1:
# 直接内部组合两个loss
ce_loss_1, pro = get_loss1(outputs[0], labels)
ce_loss = 0
else:
ce_loss = get_ce_loss(outputs[0], labels)
if Config.use_Asoftmax:
fetch_batch = labels.size(0)
if batch_cnt % (train_epoch_step // 5) == 0:
angle_loss = get_angle_loss(outputs[3],
labels[0:fetch_batch:2],
decay=0.9)
else:
angle_loss = get_angle_loss(outputs[3],
labels[0:fetch_batch:2])
loss += angle_loss
loss += ce_loss
alpha_ = 1
beta_ = 1
# gamma_ = 0.01 if Config.dataset == 'STCAR' or Config.dataset == 'AIR' else 1
gamma_ = 0.01
if Config.use_dcl:
if Config.use_focal_loss:
swap_loss = get_focal_loss(outputs[1], labels_swap) * beta_
else:
if Config.use_loss1:
swap_loss, _ = get_loss1(outputs[1],
labels_swap,
brand_prob=pro_s)
else:
swap_loss = get_ce_loss(outputs[1],
labels_swap) * beta_
loss += swap_loss
if not Config.no_loc:
law_loss = add_loss(outputs[2], swap_law) * gamma_
loss += law_loss
loss.backward()
torch.cuda.synchronize()
torch.cuda.synchronize()
if Config.use_dcl:
if Config.multi:
print(
'step: {:-8d} / {:d} loss: {:6.4f} ce_loss: {:6.4f} swap_loss: {:6.4f} '
.format(step, train_epoch_step,
loss.detach().item(),
ce_loss.detach().item(),
swap_loss.detach().item()),
flush=True)
# if Config.use_loss1:
# print(
# 'step: {:-8d} / {:d} loss: {:6.4f} ce_loss: {:6.4f} swap_loss: {:6.4f} '.format(step,train_epoch_step,loss.detach().item(),ce_loss.detach().item(),swap_loss.detach().item()),
# flush=True)
elif Config.no_loc:
print(
'step: {:-8d} / {:d} loss=ce_loss+swap_loss+law_loss: {:6.4f} = {:6.4f} + {:6.4f} '
.format(step, train_epoch_step,
loss.detach().item(),
ce_loss.detach().item(),
swap_loss.detach().item()),
flush=True)
else:
print(
'step: {:-8d} / {:d} loss=ce_loss+swap_loss+law_loss: {:6.4f} = {:6.4f} + {:6.4f} + {:6.4f} '
.format(step, train_epoch_step,
loss.detach().item(),
ce_loss.detach().item(),
swap_loss.detach().item(),
law_loss.detach().item()),
flush=True)
if Config.use_backbone:
print(
'step: {:-8d} / {:d} loss=ce_loss+swap_loss+law_loss: {:6.4f} = {:6.4f} '
.format(step, train_epoch_step,
loss.detach().item(),
ce_loss.detach().item()),
flush=True)
rec_loss.append(loss.detach().item())
train_loss_recorder.update(loss.detach().item())
# evaluation & save
if step % checkpoint == 0:
rec_loss = []
print(32 * '-', flush=True)
print(
'step: {:d} / {:d} global_step: {:8.2f} train_epoch: {:04d} rec_train_loss: {:6.4f}'
.format(step, train_epoch_step,
1.0 * step / train_epoch_step, epoch,
train_loss_recorder.get_val()),
flush=True)
print('current lr:%s' % exp_lr_scheduler.get_lr(), flush=True)
if Config.multi:
val_acc_s, val_acc_b, val_acc_c, val_acc_t = eval_turn(
Config, model, data_loader['val'], 'val', epoch)
is_best = val_acc_s > best_prec1
best_prec1 = max(val_acc_s, best_prec1)
filename = 'weights_%d_%d_%.4f_%.4f.pth' % (
epoch, batch_cnt, val_acc_s, val_acc_b)
save_checkpoint(model.state_dict(), is_best, save_dir,
filename)
sw.add_scalar("Train_Loss/Total_loss",
loss.detach().item(), epoch)
sw.add_scalar("Train_Loss/b_loss",
b_loss.detach().item(), epoch)
sw.add_scalar("Train_Loss/t_loss",
t_loss.detach().item(), epoch)
sw.add_scalar("Train_Loss/s_loss",
s_loss.detach().item(), epoch)
sw.add_scalar("Train_Loss/c_loss",
c_loss.detach().item(), epoch)
sw.add_scalar("Accurancy/val_acc_s", val_acc_s, epoch)
sw.add_scalar("Accurancy/val_acc_b", val_acc_b, epoch)
sw.add_scalar("Accurancy/val_acc_c", val_acc_c, epoch)
sw.add_scalar("Accurancy/val_acc_t", val_acc_t, epoch)
sw.add_scalar("learning_rate",
exp_lr_scheduler.get_lr()[1], epoch)
else:
val_acc1, val_acc2, val_acc3 = eval_turn(
Config, model, data_loader['val'], 'val', epoch)
is_best = val_acc1 > best_prec1
best_prec1 = max(val_acc1, best_prec1)
filename = 'weights_%d_%d_%.4f_%.4f.pth' % (
epoch, batch_cnt, val_acc1, val_acc3)
save_checkpoint(model.state_dict(), is_best, save_dir,
filename)
sw.add_scalar("Train_Loss", loss.detach().item(), epoch)
sw.add_scalar("Val_Accurancy", val_acc1, epoch)
sw.add_scalar("learning_rate",
exp_lr_scheduler.get_lr()[1], epoch)
torch.cuda.empty_cache()
# save only
elif step % savepoint == 0:
train_loss_recorder.update(rec_loss)
rec_loss = []
save_path = os.path.join(
save_dir, 'savepoint_weights-%d-%s.pth' % (step, dt()))
checkpoint_list.append(save_path)
if len(checkpoint_list) == 6:
os.remove(checkpoint_list[0])
del checkpoint_list[0]
torch.save(model.state_dict(), save_path)
torch.cuda.empty_cache()
def train(Config,
model,
epoch_num,
start_epoch,
optimizer,
exp_lr_scheduler,
data_loader,
save_dir,
data_ver='all',
data_size=448,
savepoint=500,
checkpoint=1000):
step = 0
eval_train_flag = False
rec_loss = []
checkpoint_list = []
train_batch_size = data_loader['train'].batch_size
train_epoch_step = data_loader['train'].__len__()
train_loss_recorder = LossRecord(train_batch_size)
bitempered_layer = BiTemperedLayer(t1=0.9, t2=1.05)
bitempered_loss = BiTemperedLoss()
add_loss = nn.L1Loss()
get_focal_loss = FocalLoss()
get_angle_loss = AngleLoss()
get_ce_loss = nn.CrossEntropyLoss()
get_l2_loss = nn.MSELoss()
for epoch in range(start_epoch, epoch_num - 1):
exp_lr_scheduler.step(epoch)
model.train(True)
save_grad = []
for batch_cnt, data in enumerate(data_loader['train']):
step += 1
loss = 0
model.train(True)
inputs, labels, img_names = data
inputs = inputs.cuda()
labels = torch.from_numpy(np.array(labels)).cuda()
optimizer.zero_grad()
outputs = model(inputs)
# ce_loss = get_ce_loss(outputs, labels)
labels_onehot = torch.zeros(outputs.shape[0],
50030).cuda().scatter_(
1, labels.unsqueeze_(1), 1)
ce_loss = acc_loss(labels_onehot, F.softmax(outputs, 1))
loss += ce_loss
loss.backward()
#torch.cuda.synchronize()
optimizer.step()
#torch.cuda.synchronize()
log_file.close()
def train(Config,
model,
epoch_num,
start_epoch,
optimizer,
exp_lr_scheduler,
data_loader,
save_dir,
data_size=448,
savepoint=5000,
checkpoint=5000):
# savepoint: save without evalution
# checkpoint: save with evaluation
eval_train_flag = False
rec_loss = []
checkpoint_list = []
train_batch_size = data_loader['train'].batch_size
train_epoch_step = data_loader['train'].__len__()
train_loss_recorder = LossRecord(train_batch_size)
if savepoint > train_epoch_step:
savepoint = 1 * train_epoch_step
checkpoint = savepoint
date_suffix = dt()
log_file = open(
os.path.join(
Config.log_folder,
'formal_log_r50_dcl_%s_%s.log' % (str(data_size), date_suffix)),
'a')
add_loss = nn.L1Loss()
get_ce_loss = nn.CrossEntropyLoss()
get_ce_sig_loss = nn.BCELoss()
get_focal_loss = FocalLoss()
get_angle_loss = AngleLoss()
step = 0
for epoch in range(start_epoch, epoch_num - 1):
exp_lr_scheduler.step(epoch)
model.train(True)
save_grad = []
for batch_cnt, data in enumerate(data_loader['train']):
step += 1
loss = 0
model.train(True)
if Config.use_backbone:
inputs, labels, img_names = data
inputs = Variable(inputs.cuda())
# labels = Variable(torch.LongTensor(np.array(labels)).cuda())
labels = Variable(torch.FloatTensor(np.array(labels)).cuda())
if Config.use_dcl:
inputs, labels, labels_swap, swap_law, law_index, img_names = data
inputs = Variable(inputs.cuda())
# print (type(labels))
# labels = Variable(torch.LongTensor(np.array(labels)).cuda())
labels = Variable(torch.FloatTensor(np.array(labels)).cuda())
####### dy modify
# labels_numpy = np.array(labels.cpu()).astype(np.uint8)
# print (labels_numpy)
labels_swap = Variable(
torch.LongTensor(np.array(labels_swap)).cuda())
swap_law = Variable(
torch.LongTensor(np.array(swap_law)).float().cuda())
optimizer.zero_grad()
if inputs.size(0) < 2 * train_batch_size:
outputs = model(inputs, inputs[0:-1:2])
else:
outputs = model(inputs, law_index)
idx_unswap = torch.tensor([0, 2, 4, 6, 8], dtype=torch.long).cuda()
unswap_label = torch.index_select(labels, dim=0, index=idx_unswap)
# print (inputs.size(0))
if Config.use_focal_loss:
ce_loss = get_focal_loss(outputs[0], labels)
else:
# ce_loss = get_ce_loss(outputs[0], labels) ### classification batach x 200
# print (outputs[0].shape)
# print (unswap_label.shape)
ce_loss = get_ce_sig_loss(
outputs[0], unswap_label) ### classification batach x 200
if Config.use_Asoftmax:
fetch_batch = labels.size(0)
if batch_cnt % (train_epoch_step // 5) == 0:
angle_loss = get_angle_loss(outputs[3],
labels[0:fetch_batch:2],
decay=0.9)
else:
angle_loss = get_angle_loss(outputs[3],
labels[0:fetch_batch:2])
loss += angle_loss
alpha_ = 1
loss += ce_loss * alpha_
beta_ = 0.1
gamma_ = 0.01 if Config.dataset == 'STCAR' or Config.dataset == 'AIR' else 1
if Config.use_dcl:
swap_loss = get_ce_loss(
outputs[1], labels_swap
) * beta_ ### adverisal classification batach x 2
loss += swap_loss ####### 0.692 * 0.1 = 0.0692
law_loss = add_loss(
outputs[2], swap_law
) * gamma_ ### mask L1Loss batach x 49 L1 Loss 主要用来计算 input x 和 target y 的逐元素间差值的平均绝对值.
loss += law_loss ####### 0.0683 * 1 = 0.0683
loss.backward()
torch.cuda.synchronize()
optimizer.step()
torch.cuda.synchronize()
if Config.use_dcl:
print(
'epoch:{:d}, globalstep: {:-8d}, {:d} / {:d} \n loss=ce_l+swap_l+law_l: {:6.4f} = {:6.4f} + {:6.4f} + {:6.4f} '
.format(epoch, step, batch_cnt, train_epoch_step,
loss.detach().item(),
ce_loss.detach().item(),
swap_loss.detach().item(),
law_loss.detach().item()),
flush=True)
if Config.use_backbone:
print(
'step: {:-8d} / {:d} loss=ce_loss+swap_loss+law_loss: {:6.4f} = {:6.4f} '
.format(step, train_epoch_step,
loss.detach().item(),
ce_loss.detach().item()),
flush=True)
rec_loss.append(loss.detach().item())
train_loss_recorder.update(loss.detach().item())
# evaluation & save
if step % checkpoint == 0:
rec_loss = []
print(32 * '-', flush=True)
print(
'step: {:d} / {:d} global_step: {:8.2f} train_epoch: {:04d} rec_train_loss: {:6.4f}'
.format(step, train_epoch_step,
1.0 * step / train_epoch_step, epoch,
train_loss_recorder.get_val()),
flush=True)
print('current lr:%s' % exp_lr_scheduler.get_lr(), flush=True)
val_acc = eval_turn(Config, model, data_loader['trainval'],
'val', epoch, log_file)
# if val_acc >0.9:
# checkpoint = 500
# savepoint = 500
# save_path = os.path.join(save_dir, 'weights_%d_%d_%.4f_%.4f.pth'%(epoch, batch_cnt, val_acc1, val_acc3))
save_path = os.path.join(
save_dir,
'weights_%d_%d_%.4f.pth' % (epoch, batch_cnt, val_acc))
torch.cuda.synchronize()
torch.save(model.state_dict(), save_path)
print('saved model to %s' % (save_path), flush=True)
torch.cuda.empty_cache()
# save only
elif step % savepoint == 0:
train_loss_recorder.update(rec_loss)
rec_loss = []
save_path = os.path.join(
save_dir, 'savepoint_weights-%d-%s.pth' % (step, dt()))
checkpoint_list.append(save_path)
if len(checkpoint_list) == 6:
os.remove(checkpoint_list[0])
del checkpoint_list[0]
torch.save(model.state_dict(), save_path)
torch.cuda.empty_cache()
log_file.close()
f_score = smp.utils.functional.f_score
# Dice/F1 score - https://en.wikipedia.org/wiki/S%C3%B8rensen%E2%80%93Dice_coefficient
# IoU/Jaccard score - https://en.wikipedia.org/wiki/Jaccard_index
diceLoss = smp.utils.losses.DiceLoss(eps=1)
tv_loss = kornia.losses.TotalVariation()
tv_loss.__name__ = 'tv_loss'
mseLoss = torch.nn.MSELoss(reduction='mean')
mseLoss.__name__ = 'mse_loss'
AverageValueMeter = smp.utils.train.AverageValueMeter
from models.focal_loss import FocalLoss
focal_loss = FocalLoss()
metrics = [
diceLoss,
focal_loss,
tv_loss,
smp.utils.metrics.IoU(threshold=0.5),
smp.utils.metrics.Fscore()
]
def format_logs(logs):
str_logs = ['{} - {:.4}'.format(k, v) for k, v in logs.items()]
s = ', '.join(str_logs)
return s
def main():
# 네트워크
D = resnet_face18(use_se=False).to(device)
arcmargin = ArcMarginProduct(512, celeba.classes, s=30, m=0.5, easy_margin=False).to(device)
# pretrained 모델 불러오기
if opts.reuse:
assert os.path.isfile('pretrained/trained_D.pth')
checkpoint = torch.load('pretrained/trained_D.pth')
D.load_state_dict(checkpoint['D'])
arcmargin.load_state_dict(checkpoint['arcmargin'])
print('[*]Pretrained model loaded')
# optimizer 정의
D_optim = optim.SGD([{'params':D.parameters()}, {'params':arcmargin.parameters()}],
lr=opts.lr, weight_decay=opts.weight_decay)
# loss 정의
criterion = FocalLoss(gamma=2)
mseloss = nn.MSELoss()
# scheduler
scheduler = optim.lr_scheduler.StepLR(D_optim, step_size=opts.lr_step ,gamma=0.1)
for epoch in range(opts.num_epoch):
for i, (img, cls, bbox, landm) in enumerate(dataloader):
img = img.to(device)
cls = cls.to(device)
bbox = bbox.to(device)
landm = landm.to(device)
expanded_bbox = to_expand_bbox(bbox, landm, *img.shape[2:])
cropped_img = to_crop_resize(img, expanded_bbox).to(device)
feature = D(cropped_img)
output = arcmargin(feature, cls)
loss_D = criterion(output, cls)
D_optim.zero_grad()
loss_D.backward()
D_optim.step()
# 진행상황 출력
output = output.data.cpu().numpy()
output = np.argmax(output, axis=1)
cls = cls.data.cpu().numpy()
acc = np.mean((output==cls).astype(int)).astype(float)
print(f"[Epoch {epoch}/{opts.num_epoch}] [Batch {i*opts.batch_size}/{len(celeba)}] [D loss {loss_D:.6f}] lr {scheduler.get_lr()}")
print(f"Precision | {acc:.4f}")
print("="*65)
# 로깅
if i % (len(dataloader) // opts.logstep_per_epoch) == 0 and opts.logging:
logger.write('epoch',epoch, 'iter',i, 'loss_d',loss_D, 'acc', acc)
torch.save({
'D' : D.state_dict(),
'arcmargin' : arcmargin.state_dict()
},
'pretrained/trained_D_arc.pth')
scheduler.step()