def train():
net.train()
epoch = 0 + args.resume_epoch
print('Loading Dataset...')
dataset = WiderFaceDetection(training_dataset, preproc(img_dim, rgb_mean))
testset = testDataset(test_dataset, testpreproc(img_dim, rgb_mean))
epoch_size = math.ceil(len(dataset) / batch_size)
max_iter = max_epoch * epoch_size
stepvalues = (cfg['decay1'] * epoch_size, cfg['decay2'] * epoch_size)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(data.DataLoader(dataset, batch_size, shuffle=True, num_workers=num_workers,
collate_fn=detection_collate, pin_memory=True))
if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0 and epoch > cfg['decay1']):
# loss = test(testset)
# net.train()
torch.save(net.state_dict(), save_folder + cfg['save_name'] + '_epoch_' + str(epoch) + '.pth')
# print('Epoch:{}/{} || loss: {}'.format(epoch, max_epoch, loss))
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(optimizer, gamma, epoch, step_index, iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
images = images.cuda()
targets = [anno.cuda() for anno in targets]
# forward
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c, loss_landm = criterion(out, priors, targets)
loss = cfg['loc_weight'] * loss_l + loss_c + loss_landm
loss.backward()
optimizer.step()
load_t1 = time.time()
batch_time = load_t1 - load_t0
eta = int(batch_time * (max_iter - iteration))
print(
'Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || Loc: {:.4f} Cla: {:.4f} Landm: {:.4f} || LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}'
.format(epoch, max_epoch, (iteration % epoch_size) + 1,
epoch_size, iteration + 1, max_iter, loss_l.item(), loss_c.item(), loss_landm.item(), lr,
batch_time, str(datetime.timedelta(seconds=eta))))
torch.save(net.state_dict(), save_folder + cfg['save_name'] + '_Final.pth')
def main():
if args.resume_net is not None and os.path.isfile(args.resume_net):
checkpoint = torch.load(args.resume_net, map_location="cpu")
cfg = checkpoint["config"]
else:
checkpoint = None
if args.network == "mobilenet0.25":
cfg = cfg_mnet
elif args.network == "resnet50":
cfg = cfg_re50
cfg, net = initialize_network(cfg, checkpoint)
torch.backends.cudnn.benchmark = True
optimizer = torch.optim.SGD(
net.parameters(), lr=initial_lr,
momentum=args.momentum, weight_decay=args.weight_decay,
)
criterion = MultiBoxLoss(2, 0.35, True, 0, True, 7, 0.35, False)
dataset = WiderFaceDetection(training_dataset, preproc(img_dim, RGB_MEAN))
dataloader = torch.utils.data.DataLoader(
dataset, batch_size, shuffle=True,
num_workers=args.num_workers, collate_fn=dataset.collate,
)
training_loop(net, optimizer, criterion, dataloader, cfg)
def train():
net.train()
epoch = 0 + args.resume_epoch
pickleFileSaverName=input("Enter the name to save loss data : ")
print('Loading Dataset...')
dataset = WiderFaceDetection( training_dataset,preproc(img_dim, rgb_mean))
print('Loading Val Dataset...')
val_data = WiderFaceDetection(validation_dataset,preproc(img_dim, rgb_mean))
dataset_ = data.DataLoader(val_data,batch_size, shuffle=True, num_workers=num_workers, collate_fn=detection_collate)
epoch_size = math.ceil(len(dataset) / batch_size)
max_iter = max_epoch * epoch_size
stepvalues = (cfg['decay1'] * epoch_size, cfg['decay2'] * epoch_size)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
epoch_loss_train = 0.0
lossCollector=[]
lossrn=2000#just a random value greater than 155 to move on with the first epoch
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(data.DataLoader(dataset, batch_size, shuffle=True, num_workers=num_workers, collate_fn=detection_collate))
if (epoch % 2 == 0 and epoch > 0) or (epoch % 5 == 0 and epoch > cfg['decay1']):
torch.save(net.state_dict(), save_folder + cfg['name']+ '_epoch_' + str(epoch) + '_noGrad_FT_Adam_WC1.pth')
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate_by_neels(optimizer,lossrn)
# load train data
images, targets = next(batch_iterator)
images = images.cuda()
targets = [anno.cuda() for anno in targets]
# forward
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c, loss_landm = criterion(out, priors, targets)
loss = cfg['loc_weight'] * loss_l + loss_c + loss_landm
loss.backward()
optimizer.step()
load_t1 = time.time()
batch_time = load_t1 - load_t0
eta = int(batch_time * (max_iter - iteration))
epoch_loss_train += loss.item()
print('Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || Loc: {:.4f} Cla: {:.4f} Landm: {:.4f} || LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}'
.format(epoch, max_epoch, (iteration % epoch_size) + 1,
epoch_size, iteration + 1, max_iter, loss_l.item(), loss_c.item(), loss_landm.item(), lr, batch_time, str(datetime.timedelta(seconds=eta))))
# calculate validation loss after each epoch
if iteration % epoch_size == 0:
# print('Training loss per image for Epoch {} : {}'.format(epoch,epoch_loss_train/len(dataset)))
print('Training loss for Epoch {} : {}'.format(epoch,epoch_loss_train))
valLoss=train_eval(net,dataset_,batch_size,epoch)
lossrn=valLoss
lossCollector.append({"Epoch":epoch,"TrainLoss":epoch_loss_train,"ValLoss":valLoss})
#saving the losses data per epoch
picklefile=open("./lossData/"+pickleFileSaverName+"_{}.pickle".format(epoch),"wb")
pickle.dump(lossCollector,picklefile)
picklefile.close()
#saving is complete
epoch_loss_train = 0.0
torch.save(net.state_dict(), save_folder + cfg['name'] + '_Finally_FT_Adam_WC1.pth')
# torch.save(net.state_dict(), save_folder + 'Final_Retinaface.pth')
#saving the data for losses
picklefile=open("./lossData/"+pickleFileSaverName+".pickle" ,"wb")
pickle.dump(lossCollector,picklefile)
picklefile.close()
def main():
args = get_args()
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
cfg = None
if args.network == "mobile0.25":
cfg = cfg_mnet
elif args.network == "resnet50":
cfg = cfg_re50
rgb_mean = (104, 117, 123) # bgr order
num_classes = 2
img_dim = cfg["image_size"]
num_gpu = cfg["ngpu"]
batch_size = cfg["batch_size"]
max_epoch = cfg["epoch"]
gpu_train = cfg["gpu_train"]
num_workers = args.num_workers
momentum = args.momentum
weight_decay = args.weight_decay
initial_lr = args.lr
gamma = args.gamma
training_dataset = args.training_dataset
save_folder = args.save_folder
net = RetinaFace(cfg=cfg)
print("Printing net...")
print(net)
if args.resume_net is not None:
print("Loading resume network...")
state_dict = torch.load(args.resume_net)
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == "module.":
name = k[7:] # remove `module.`
else:
name = k
new_state_dict[name] = v
net.load_state_dict(new_state_dict)
if num_gpu > 1 and gpu_train:
net = torch.nn.DataParallel(net).cuda()
else:
net = net.cuda()
cudnn.benchmark = True
optimizer = optim.SGD(net.parameters(),
lr=initial_lr,
momentum=momentum,
weight_decay=weight_decay)
criterion = MultiBoxLoss(num_classes, 0.35, True, 0, True, 7, 0.35, False)
priorbox = PriorBox(cfg, image_size=(img_dim, img_dim))
with torch.no_grad():
priors = priorbox.forward()
priors = priors.cuda()
net.train()
epoch = 0 + args.resume_epoch
print("Loading Dataset...")
dataset = WiderFaceDetection(training_dataset, preproc(img_dim, rgb_mean))
epoch_size = math.ceil(len(dataset) / batch_size)
max_iter = max_epoch * epoch_size
stepvalues = (cfg["decay1"] * epoch_size, cfg["decay2"] * epoch_size)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(
data.DataLoader(dataset,
batch_size,
shuffle=True,
num_workers=num_workers,
collate_fn=detection_collate))
if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0
and epoch > cfg["decay1"]):
torch.save(
net.state_dict(), save_folder + cfg["name"] + "_epoch_" +
str(epoch) + ".pth")
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(initial_lr, optimizer, gamma, epoch,
step_index, iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
images = images.cuda()
targets = [anno.cuda() for anno in targets]
# forward
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c, loss_landm = criterion(out, priors, targets)
loss = cfg["loc_weight"] * loss_l + loss_c + loss_landm
loss.backward()
optimizer.step()
load_t1 = time.time()
batch_time = load_t1 - load_t0
eta = int(batch_time * (max_iter - iteration))
print(
"Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || Loc: {:.4f} Cla: {:.4f} Landm: {:.4f} "
"|| LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}".format(
epoch,
max_epoch,
(iteration % epoch_size) + 1,
epoch_size,
iteration + 1,
max_iter,
loss_l.item(),
loss_c.item(),
loss_landm.item(),
lr,
batch_time,
str(datetime.timedelta(seconds=eta)),
))
torch.save(net.state_dict(), save_folder + cfg["name"] + "_Final.pth")
def train():
net.train()
epoch = 0 + args.resume_epoch
print('Loading Dataset...')
writer = SummaryWriter(logdir=args.log_folder)
dataset = WiderFaceDetection(training_dataset,
preproc(img_dim, rgb_mean, rgb_std))
print("batch_size", batch_size)
epoch_size = math.ceil(len(dataset) / batch_size)
max_iter = max_epoch * epoch_size
stepvalues = (cfg['decay1'] * epoch_size, cfg['decay2'] * epoch_size)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(
data.DataLoader(dataset,
batch_size,
shuffle=True,
num_workers=num_workers,
collate_fn=detection_collate))
if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0
and epoch > cfg['decay1']):
torch.save(
net.state_dict(), save_folder + cfg['name'] + '_epoch_' +
str(epoch) + '.pth')
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(optimizer, gamma, epoch, step_index,
iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
images = images.cuda()
targets = [anno.cuda() for anno in targets]
# forward
out = net(images)
# backprop
optimizer.zero_grad()
if cfg['gender']:
loss_l, loss_c, loss_landm, loss_gender = criterion(
out, priors, targets)
loss = cfg[
'loc_weight'] * loss_l + loss_c + loss_landm + loss_gender
else:
loss_l, loss_c, loss_landm = criterion(out, priors, targets)
loss = cfg['loc_weight'] * loss_l + loss_c + loss_landm
loss.backward()
optimizer.step()
load_t1 = time.time()
batch_time = load_t1 - load_t0
eta = int(batch_time * (max_iter - iteration))
if iteration % 100 == 0:
writer.add_scalar('training/loss', loss.item(), iteration)
writer.add_scalar('training/loss_localization', loss_l.item(),
iteration)
writer.add_scalar('training/loss_confidence', loss_c.item(),
iteration)
writer.add_scalar('training/loss_landmark', loss_landm.item(),
iteration)
if cfg['gender']:
writer.add_scalar('training/loss_gender', loss_gender.item(),
iteration)
for param_group in optimizer.param_groups:
writer.add_scalar('LearningRate', param_group['lr'], iteration)
if cfg['gender']:
print(
'Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || Loc: {:.4f} Cla: {:.4f} Landm: {:.4f} Gender: {:.4f} || LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}'
.format(epoch, max_epoch, (iteration % epoch_size) + 1,
epoch_size, iteration + 1, max_iter, loss_l.item(),
loss_c.item(), loss_landm.item(), loss_gender.item(),
lr, batch_time, str(datetime.timedelta(seconds=eta))))
else:
print(
'Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || Loc: {:.4f} Cla: {:.4f} Landm: {:.4f} || LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}'
.format(epoch, max_epoch, (iteration % epoch_size) + 1,
epoch_size, iteration + 1, max_iter, loss_l.item(),
loss_c.item(), loss_landm.item(), lr, batch_time,
str(datetime.timedelta(seconds=eta))))
torch.save(net.state_dict(), save_folder + cfg['name'] + '_Final.pth')
def main():
global args
global minmum_loss
args.gpu = 0
args.world_size = 1
if args.distributed:
args.gpu = args.local_rank % torch.cuda.device_count()
torch.cuda.set_device(args.gpu)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.total_batch_size = args.world_size * args.batch_size
# build dsfd network
print("Building net...")
model = RetinaFace(cfg=cfg)
print("Printing net...")
# for multi gpu
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(model)
model = model.cuda()
# optimizer and loss function
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
criterion = MultiBoxLoss(cfg['num_classes'], 0.35, True, 0, True, 7, 0.35,
False)
## dataset
print("loading dataset")
train_dataset = WiderFaceDetection(
args.training_dataset, preproc(cfg['image_size'], cfg['rgb_mean']))
train_loader = data.DataLoader(train_dataset,
args.batch_size,
num_workers=args.num_workers,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(
args.resume,
map_location=lambda storage, loc: storage.cuda(args.gpu))
args.start_epoch = checkpoint['epoch']
minmum_loss = checkpoint['minmum_loss']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
print('Using the specified args:')
print(args)
# load PriorBox
print("Load priorbox")
with torch.no_grad():
priorbox = PriorBox(cfg=cfg,
image_size=(cfg['image_size'], cfg['image_size']))
priors = priorbox.forward()
priors = priors.cuda()
print("start traing")
for epoch in range(args.start_epoch, args.epochs):
# train for one epoch
train_loss = train(train_loader, model, priors, criterion, optimizer,
epoch)
if args.local_rank == 0:
is_best = train_loss < minmum_loss
minmum_loss = min(train_loss, minmum_loss)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': minmum_loss,
'optimizer': optimizer.state_dict(),
}, is_best, epoch)
def train():
net.train()
epoch = 0 + args.resume_epoch
print('Loading Dataset...')
dataset = WiderFaceDetection(training_dataset, preproc(img_dim, rgb_mean))
val_dataset = WiderFaceDetection(training_dataset, preproc(img_dim, rgb_mean), phase='val')
epoch_size = math.ceil(len(dataset) / batch_size)
max_iter = max_epoch * epoch_size
stepvalues = (cfg['decay1'] * epoch_size, cfg['decay2'] * epoch_size)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
for iteration in tqdm(range(start_iter, max_iter)):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(data.DataLoader(
dataset, batch_size, shuffle=True, num_workers=num_workers, pin_memory=True, collate_fn=detection_collate))
if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0 and epoch > cfg['decay1']):
ckpt_path = save_folder + \
cfg['name'] + '_epoch_' + str(epoch) + '.pth'
torch.save(net.state_dict(), ckpt_path)
# print('Validating....')
# net.eval()
# txt_save_folder = val_to_text()
# aps = evaluation(
# txt_save_folder, 'widerface_evaluate/ground_truth/')
# writer.add_scalar('Val AP/easy_AP', aps[0], epoch)
# writer.add_scalar('Val AP/medium_AP', aps[1], epoch)
# writer.add_scalar('Val AP/hard_AP', aps[2], epoch)
val(val_dataset, epoch)
net.train()
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(
optimizer, gamma, epoch, step_index, iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
images = images.cuda()
targets = [anno.cuda() for anno in targets]
# print('targets', len(targets))
# print('images', images.shape)
# for i in range(len(targets)):
# print(targets[i].shape)
# forward
out = net(images)
# dad_out = dad_net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c, loss_landm = criterion(out, priors, targets)
# loss_dis = distillation_loss(out, dad_out)
# loss_l, loss_c, loss_landm = 0 * loss_dis, 0 * loss_dis, 0 * loss_dis
loss = loss_l + loss_c + loss_landm
# loss += 5 * loss_dis
loss.backward()
optimizer.step()
load_t1 = time.time()
batch_time = load_t1 - load_t0
eta = int(batch_time * (max_iter - iteration))
if iteration % 100 == 0:
writer.add_scalar('Train loss/loss_total', loss.item(), iteration)
writer.add_scalar('Train loss/loss_bbox', loss_l.item(), iteration)
writer.add_scalar('Train loss/loss_class',
loss_c.item(), iteration)
writer.add_scalar('Train loss/loss_lm',
loss_landm.item(), iteration)
# writer.add_scalar('Train loss/loss_dis',
# loss_dis.item(), iteration)
writer.add_scalar('LR', lr, iteration)
if iteration % 1000 == 0:
print('Epoch:{}/{} || Iter: {}/{} || Loss_bbox: {:.4f} Loss_class: {:.4f} Loss_lm: {:.4f} || Batchtime: {:.4f} s || ETA: {}'
.format(epoch, max_epoch,
iteration + 1, max_iter,
loss_l.item(),
loss_c.item(),
loss_landm.item(),
batch_time,
str(datetime.timedelta(seconds=eta))))
torch.save(net.state_dict(), save_folder + cfg['name'] + '_Final.pth')
def train():
net = RetinaFace(cfg=cfg)
logger.info("Printing net...")
logger.info(net)
if args.resume_net is not None:
logger.info('Loading resume network...')
state_dict = torch.load(args.resume_net)
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
new_state_dict[name] = v
net.load_state_dict(new_state_dict)
if num_gpu > 1 and gpu_train:
net = torch.nn.DataParallel(net).cuda()
else:
net = net.cuda()
cudnn.benchmark = True
priorbox = PriorBox(cfg, image_size=(img_dim, img_dim))
with torch.no_grad():
priors = priorbox.forward()
priors = priors.cuda()
net.train()
epoch = 0 + args.resume_epoch
logger.info('Loading Dataset...')
trainset = WiderFaceDetection(training_dataset, preproc=train_preproc(img_dim, rgb_mean), mode='train')
validset = WiderFaceDetection(training_dataset, preproc=valid_preproc(img_dim, rgb_mean), mode='valid')
# trainset = WiderFaceDetection(training_dataset, transformers=train_transformers(img_dim), mode='train')
# validset = WiderFaceDetection(training_dataset, transformers=valid_transformers(img_dim), mode='valid')
trainloader = data.DataLoader(trainset, batch_size, shuffle=True, num_workers=num_workers, collate_fn=detection_collate)
validloader = data.DataLoader(validset, batch_size, shuffle=True, num_workers=num_workers, collate_fn=detection_collate)
logger.info(f'Totally {len(trainset)} training samples and {len(validset)} validating samples.')
epoch_size = math.ceil(len(trainset) / batch_size)
max_iter = max_epoch * epoch_size
logger.info(f'max_epoch: {max_epoch:d} epoch_size: {epoch_size:d}, max_iter: {max_iter:d}')
# optimizer = optim.SGD(net.parameters(), lr=initial_lr, momentum=momentum, weight_decay=weight_decay)
optimizer = optim.Adam(net.parameters(), lr=initial_lr, weight_decay=weight_decay)
scheduler = _utils.get_linear_schedule_with_warmup(optimizer, int(0.1 * max_iter), max_iter)
criterion = MultiBoxLoss(num_classes, 0.35, True, 0, True, 7, 0.35, False)
stepvalues = (cfg['decay1'] * epoch_size, cfg['decay2'] * epoch_size)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
best_loss_val = float('inf')
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# create batch iterator
# batch_iterator = iter(tqdm(trainloader, total=len(trainloader)))
batch_iterator = iter(trainloader)
# if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0 and epoch > cfg['decay1']):
# torch.save(net.state_dict(), save_folder + cfg['name']+ '_epoch_' + str(epoch) + '.pth')
epoch += 1
torch.cuda.empty_cache()
if (valid_steps > 0) and (iteration > 0) and (iteration % valid_steps == 0):
net.eval()
# validation
loss_l_val = 0.
loss_c_val = 0.
loss_landm_val = 0.
loss_val = 0.
# for val_no, (images, targets) in tqdm(enumerate(validloader), total=len(validloader)):
for val_no, (images, targets) in enumerate(validloader):
# load data
images = images.cuda()
targets = [anno.cuda() for anno in targets]
# forward
with torch.no_grad():
out = net(images)
loss_l, loss_c, loss_landm = criterion(out, priors, targets)
loss = cfg['loc_weight'] * loss_l + loss_c + loss_landm
loss_l_val += loss_l.item()
loss_c_val += loss_c.item()
loss_landm_val += loss_landm.item()
loss_val += loss.item()
loss_l_val /= len(validloader)
loss_c_val /= len(validloader)
loss_landm_val /= len(validloader)
loss_val /= len(validloader)
logger.info('[Validating] Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || Total: {:.4f} Loc: {:.4f} Cla: {:.4f} Landm: {:.4f}'
.format(epoch, max_epoch, (iteration % epoch_size) + 1,
epoch_size, iteration + 1, max_iter,
loss_val, loss_l_val, loss_c_val, loss_landm_val))
if loss_val < best_loss_val:
best_loss_val = loss_val
pth = os.path.join(save_folder, cfg['name'] + '_iter_' + str(iteration) + f'_{loss_val:.4f}_' + '.pth')
torch.save(net.state_dict(), pth)
logger.info(f'Best validating loss: {best_loss_val:.4f}, model saved as {pth:s})')
net.train()
load_t0 = time.time()
# if iteration in stepvalues:
# step_index += 1
# lr = adjust_learning_rate(optimizer, gamma, epoch, step_index, iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
images = images.cuda()
targets = [anno.cuda() for anno in targets]
# forward
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c, loss_landm = criterion(out, priors, targets)
loss = cfg['loc_weight'] * loss_l + loss_c + loss_landm
loss.backward()
optimizer.step()
scheduler.step()
load_t1 = time.time()
batch_time = load_t1 - load_t0
eta = int(batch_time * (max_iter - iteration))
if iteration % verbose_steps == 0:
logger.info('[Training] Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || Total: {:.4f} Loc: {:.4f} Cla: {:.4f} Landm: {:.4f} || LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}'
.format(epoch, max_epoch, (iteration % epoch_size) + 1,
epoch_size, iteration + 1, max_iter,
loss.item(), loss_l.item(), loss_c.item(), loss_landm.item(),
scheduler.get_last_lr()[-1], batch_time, str(datetime.timedelta(seconds=eta))))
def train(net, cfg, resume_epoch):
torch.set_grad_enabled(True)
rgb_mean = (104, 117, 123) # bgr order
num_classes = 2
img_dim = cfg['image_size']
batch_size = cfg['batch_size']
max_epoch = cfg['epoch']
# gpu_train = cfg['gpu_train']
num_workers = 4
momentum = 0.9
weight_decay = 5e-4
initial_lr = 1e-3
gamma = 0.1
training_dataset = './data/widerface/train/label.txt'
# cudnn.benchmark = True
optimizer = optim.SGD(net.parameters(),
lr=initial_lr,
momentum=momentum,
weight_decay=weight_decay)
criterion = MultiBoxLoss(num_classes, 0.35, True, 0, True, 7, 0.35, False)
priorbox = PriorBox(cfg, image_size=(img_dim, img_dim))
with torch.no_grad():
priors = priorbox.forward()
priors = priors.cuda()
net.train()
epoch = resume_epoch
print('Loading Dataset...')
dataset = WiderFaceDetection(training_dataset, preproc(img_dim, rgb_mean))
epoch_size = math.ceil(len(dataset) / batch_size)
max_iter = max_epoch * epoch_size
stepvalues = (cfg['decay1'] * epoch_size, cfg['decay2'] * epoch_size)
step_index = 0
if resume_epoch > 0:
start_iter = resume_epoch * epoch_size
else:
start_iter = 0
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(
data.DataLoader(dataset,
batch_size,
shuffle=True,
num_workers=num_workers,
collate_fn=detection_collate))
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(optimizer, gamma, epoch, step_index,
iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
images = images.to(device)
targets = [anno.to(device) for anno in targets]
# forward
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c, loss_landm = criterion(out, priors, targets)
loss = cfg['loc_weight'] * loss_l + loss_c + loss_landm
loss.backward()
optimizer.step()
load_t1 = time.time()
batch_time = load_t1 - load_t0
eta = int(batch_time * (max_iter - iteration))
print(
'Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || Loc: {:.4f} Cla: {:.4f} Landm: {:.4f} || LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}'
.format(epoch, max_epoch, (iteration % epoch_size) + 1, epoch_size,
iteration + 1, max_iter, loss_l.item(), loss_c.item(),
loss_landm.item(), lr, batch_time,
str(datetime.timedelta(seconds=eta))))
cls_mask_list.append(cls_mask.detach())
loc_mask_list.append(loc_mask.detach())
return loc_mask_list, cls_mask_list
if __name__ == '__main__':
num_classes = 2
fdm = FeatureDistillMask(num_classes, 3, 2)
training_dataset = './data/widerface/train/label.txt'
rgb_mean = (104, 117, 123) # bgr order
img_dim = 640
batch_size = 4
dataset = WiderFaceDetection(training_dataset, preproc(img_dim, rgb_mean))
batch_iterator = iter(
data.DataLoader(dataset,
batch_size,
shuffle=True,
num_workers=1,
collate_fn=detection_collate))
net = RetinaFace().cuda()
priorbox = PriorBox(cfg, image_size=(img_dim, img_dim))
with torch.no_grad():
priors, priors_by_layer = priorbox.forward()
priors = priors.cuda()
for iteration in range(0, 1):
def train():
net.train()
epoch = 0 + args.resume_epoch
print('Loading Dataset...')
dataset = WiderFaceDetection(training_dataset,
preproc(img_dim, rgb_mean),
cutmix_prob=0.6)
epoch_size = math.ceil(len(dataset) / batch_size)
if args.optim == 'SGD':
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=5e-4)
elif args.optim == 'Adam':
optimizer = optim.Adam(net.parameters(), lr=args.lr, weight_decay=1e-4)
else:
raise ValueError('wrong opitm')
if args.lr_type == 'plateau':
lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
mode='max',
factor=0.5,
patience=0,
min_lr=1e-4)
elif args.lr_type == 'cos':
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer,
T_max=10)
else:
lr_scheduler = None
#from copy import deepcopy
torch.multiprocessing.set_sharing_strategy('file_system')
#train_loader = data.DataLoader(dataset, batch_size, shuffle=True, num_workers=num_workers, collate_fn=detection_collate, drop_last=True)
#for images, targets in train_loader:
# print('batch')
#print('done')
train_loader = data.DataLoader(dataset,
batch_size,
shuffle=True,
num_workers=num_workers,
collate_fn=detection_collate)
best_metrics = [0.] * 4
if args.resume_net is not None:
aps = validate(args.resume_net, args.network)
print('val metrics:', aps)
best_metrics = aps
#for iteration in range(start_iter, max_iter):
for epoch in range(max_epoch):
load_t0 = time.time()
for batch_idx, (images, targets) in enumerate(train_loader):
images = images.cuda()
targets = [anno.cuda() for anno in targets]
#mixup
if False: #random.random() < 0.4 and epoch >= 0:
shuffle_indices = torch.randperm(images.size(0))
indices = torch.arange(images.size(0))
lam = np.clip(np.random.beta(1.0, 1.0), 0.35, 0.65)
images = lam * images + (1 - lam) * images[shuffle_indices, :]
mix_targets = []
for i, si in zip(indices, shuffle_indices):
#print(targets[i.item()].size())
if i.item() == si.item():
target = targets[i.item()]
else:
target = torch.cat(
[targets[i.item()], targets[si.item()]], 0)
mix_targets.append(target)
#targets = torch.cat(mix_targets, 0)
#print('mix:', len(targets), targets[0].size())
else:
#print('norm:', len(targets), targets[0].size())
pass
# forward
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c, loss_landm = criterion(out, priors, targets)
loss = cfg['loc_weight'] * loss_l + loss_c + loss_landm
loss.backward()
optimizer.step()
load_t1 = time.time()
#batch_time = load_t1 - load_t0
lr = get_lrs(optimizer)[0]
print(
'Epoch:{}/{} || Iter: {}/{} || Loc: {:.4f} Cla: {:.4f} Landm: {:.4f} || LR: {} || {:.2f} min'
.format(epoch, max_epoch, batch_idx, epoch_size, loss_l.item(),
loss_c.item(), loss_landm.item(), lr,
(load_t1 - load_t0) / 60),
end='\r')
#epoch done
print("")
torch.save(net.state_dict(), save_folder + cfg['name'] + '_latest.pth')
lr_scheduler.step()
# evaluate
if epoch % 3 == 0:
aps = validate(save_folder + cfg['name'] + '_latest.pth',
args.network)
print('val metrics:', aps)
print('cur best metrics:', best_metrics)
if aps[-1] > best_metrics[-1]:
best_metrics = aps
torch.save(net.state_dict(),
save_folder + cfg['name'] + '_best.pth')
def train():
net.train()
teacher_net.eval()
for param in teacher_net.parameters():
param.requires_grad = False
epoch = 0 + args.resume_epoch
print('Loading Dataset...')
dataset = WiderFaceDetection( training_dataset,preproc(img_dim, rgb_mean))
epoch_size = math.ceil(len(dataset) / batch_size)
max_iter = max_epoch * epoch_size
stepvalues = (55 * epoch_size, 68 * epoch_size)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(data.DataLoader(dataset, batch_size, shuffle=True, num_workers=num_workers, collate_fn=detection_collate))
if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0 and epoch > 200):
torch.save(net.state_dict(), save_folder + 'Retinaface_epoch_' + str(epoch) + '.pth')
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(optimizer, gamma, epoch, step_index, iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
images = images.to(device)
targets = [anno.to(device) for anno in targets]
# forward
out = net(images)
teacher_out = teacher_net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c, loss_landm, pos, neg = criterion(out, priors, targets)
#distillation
feat_loss, prob_loss = tdkd(out, teacher_out, priors_by_layer, priors, targets, pos, neg)
d_feat = 1.0
scale = loss_l.detach() / feat_loss.detach()
d_prob = 50
if mode == 'teacher' or mode == 'student':
d_feat = 0
d_prob = 0
scale = 0
loss = cfg['loc_weight'] * loss_l + loss_c + loss_landm + d_feat * scale * feat_loss + d_prob * prob_loss
loss.backward()
optimizer.step()
load_t1 = time.time()
batch_time = load_t1 - load_t0
eta = int(batch_time * (max_iter - iteration))
print('Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || Loc: {:.4f} Cla: {:.4f} Landm: {:.4f} Dfeat: {:.4f} Dprob: {:.4f}|| LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}'
.format(epoch, max_epoch, (iteration % epoch_size) + 1,
epoch_size, iteration + 1, max_iter, loss_l.item(), loss_c.item(), loss_landm.item(), feat_loss.item() * d_feat, prob_loss.item() * d_prob, lr, batch_time, str(datetime.timedelta(seconds=eta))))
torch.save(net.state_dict(), save_folder + 'Final_Retinaface.pth')
boxes = annotations[:, :4]
landmarks = annotations[:, 4:14]
boxes[:, 0::2] *= width
boxes[:, 1::2] *= height
landmarks[:, 0::2] *= width
landmarks[:, 1::2] *= height
landmarks = landmarks.reshape(-1, 5, 2)
image = image * 128.0 + 127.5
image = image.numpy().transpose(1, 2, 0).astype(np.uint8)
return image, boxes, landmarks
_utils.seed_everything(99)
training_dataset = '../data/widerface/WIDER_train/label.txt'
trainset = WiderFaceDetection(training_dataset,
transformers=transformers(480),
mode='train')
# for i in range(len(trainset)):
# try:
# image, annotations = trainset[i]
# print(i)
# image, boxes, landmarks = process(image, annotations)
# image = da.visualize(image, boxes, landmarks)
# cv2.imshow('', image)
# cv2.imwrite('/home/louishsu/Desktop/%d.jpg' % i, image)
# cv2.waitKey(0)
# except Exception as e:
# print(e)
# continue
optimizer = optim.SGD(net.parameters(),
lr=initial_lr,
momentum=momentum,
weight_decay=weight_decay)
criterion = MultiBoxLoss(num_classes, 0.35, True, 0, True, 7, 0.35, False)
priorbox = PriorBox(cfg, image_size=(img_dim, img_dim))
with torch.no_grad():
priors = priorbox.forward()
priors = priors.cuda()
train_path = '/opt/tiger/jupter_web/data/face_detector/train_new_visible.txt'
val_path = '/opt/tiger/jupter_web/data/face_detector/val_new_visible.txt'
val_dataset = WiderFaceDetection(training_dataset, val_path, 'val',
preproc(img_dim, rgb_mean))
def val():
net.eval()
val_loader = data.DataLoader(val_dataset,
batch_size,
shuffle=True,
num_workers=16,
collate_fn=detection_collate)
loss_l_all, loss_c_all, loss_landm_all, loss_vis_all = 0, 0, 0, 0
idx = 0
for img, target in val_loader:
img = img.cuda()
target = [anno.cuda() for anno in target]
def train():
net.train()
epoch = 0 + args.resume_epoch
print('Loading Dataset...')
# INITIALIZE DATASET
dataset = WiderFaceDetection(training_dataset,preproc(img_dim, rgb_mean))
epoch_size = math.ceil(len(dataset) / batch_size)
max_iter = max_epoch * epoch_size
stepvalues = (cfg['decay1'] * epoch_size, cfg['decay2'] * epoch_size)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(data.DataLoader(dataset, batch_size, shuffle=True, num_workers=num_workers, collate_fn=detection_collate))
if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0 and epoch > cfg['decay1']):
torch.save(net.state_dict(), save_folder + cfg['name']+ '_epoch_' + str(epoch) + '.pth')
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(optimizer, gamma, epoch, step_index, iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
images = images.cuda()
targets = [anno.cuda() for anno in targets]
# len(targets) == 32
# targets[0].shape == torch.Size([num_of_faces_in_this_image, 15]), 15 == 4(bbox) + 5*2(landmark) + 1 (have_landmark)
# forward
# torch.Size([32, 3, 640, 640]) # batch 32, image shape 3,640,640
out = net(images)
# (([32, 16800, 4]), ([32, 16800, 2]), ([32, 16800, 10])) # batch 32, 16800 anchors, for each anchor predict bbox 4, class 2, landmark 10
# backprop
optimizer.zero_grad()
loss_l, loss_c, loss_landm = criterion(out, priors, targets)
loss = cfg['loc_weight'] * loss_l + loss_c + loss_landm
loss.backward()
optimizer.step()
load_t1 = time.time()
batch_time = load_t1 - load_t0
eta = int(batch_time * (max_iter - iteration))
if ((iteration % epoch_size) + 1) % 400 == 0 or iteration < 10:
print('Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || Loc: {:.4f} Cla: {:.4f} Landm: {:.4f} || LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}'
.format(epoch, max_epoch, (iteration % epoch_size) + 1,
epoch_size, iteration + 1, max_iter, loss_l.item(), loss_c.item(), loss_landm.item(), lr, batch_time, str(datetime.timedelta(seconds=eta))))
torch.save(net.state_dict(), save_folder + cfg['name'] + '_Final.pth')
def train():
epoch = 0 + args.resume_epoch
print('Loading Dataset...')
pattern = 'train'
# print('rgb {}'.format(rgb_mean))
dataset = WiderFaceDetection(training_dataset, train_path, pattern,
preproc(img_dim, rgb_mean))
epoch_size = math.ceil(len(dataset) / batch_size)
max_iter = max_epoch * epoch_size
stepvalues = (cfg['decay1'] * epoch_size, cfg['decay2'] * epoch_size)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
for iteration in range(start_iter, max_iter):
net.train()
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(
data.DataLoader(dataset,
batch_size,
shuffle=True,
num_workers=num_workers,
collate_fn=detection_collate))
if (epoch % 1 == 0 and epoch > 0) or (epoch % 5 == 0
and epoch > cfg['decay1']):
# if epoch > 1:
# shutil.rmtree('./nas_val/')
torch.save(
net.state_dict(), save_folder + cfg['name'] + '_epoch_' +
str(epoch) + '.pth')
val_loss_l, val_loss_c, val_loss_landm, val_loss_vis = val()
logging.info(
'Validation || Epoch:{} || Loc: {:.4f} || Cla: {:.4f} || Landm: {:.4f} || Vis: {:.4f}'
.format(epoch, val_loss_l, val_loss_c, val_loss_landm,
val_loss_vis))
# val_vis(net, cfg)
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(optimizer, gamma, epoch, step_index,
iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
images = images.cuda()
# print('image {}'.format(images.shape))
targets = [anno.cuda() for anno in targets]
# print('landmark {}'.format(targets[2].shape))
# forward
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c, loss_landm, loss_vis = criterion(out, priors, targets)
# print('loss vis {}'.format(loss_vis.type()))
loss = cfg[
'loc_weight'] * loss_l + 1.5 * loss_c + 1.5 * loss_landm + 0.5 * loss_vis
# # freeze parameters
# loss = loss_landm
loss.backward()
optimizer.step()
load_t1 = time.time()
batch_time = load_t1 - load_t0
eta = int(batch_time * (max_iter - iteration))
logging.basicConfig(level=logging.INFO,
format='%(asctime)s %(message)s',
datefmt='%a, %d %b %Y %H:%M:%S +0000',
filename='result.log')
logging.info(
'Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || Loc: {:.4f} Cla: {:.4f} Landm: {:.4f} || Visible: {:.4f} || LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}'
.format(epoch, max_epoch, (iteration % epoch_size) + 1, epoch_size,
iteration + 1, max_iter, loss_l.item(), loss_c.item(),
loss_landm.item(), loss_vis.item(), lr, batch_time,
str(datetime.timedelta(seconds=eta))))
torch.save(net.state_dict(), save_folder + cfg['name'] + '_Final.pth')
def train():
net.train()
epoch = 0 + args.resume_epoch
print('Loading Dataset...')
path_dict = get_path_dict(args.train_dir)
dataset = WiderFaceDetection(training_dataset, preproc(img_dim, rgb_mean),
path_dict)
epoch_size = math.ceil(len(dataset) / batch_size)
max_iter = max_epoch * epoch_size
stepvalues = (cfg['decay1'] * epoch_size, cfg['decay2'] * epoch_size)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
# Make a file to save losses
if not os.path.isdir('./results'):
os.mkdir('./results')
with open('./results/loss.txt', 'w') as f:
pass
losses = []
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(
data.DataLoader(dataset,
batch_size,
shuffle=True,
num_workers=num_workers,
collate_fn=detection_collate))
if (epoch % 5 == 0 and epoch >
0): # or (epoch % 5 == 0 and epoch > cfg['decay1']):
torch.save(
net.state_dict(), save_folder + cfg['name'] + '_epoch_' +
str(epoch) + '.pth')
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(optimizer, gamma, epoch, step_index,
iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
images = images.cuda()
targets = [anno.cuda() for anno in targets]
# forward
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c, loss_landm = criterion(out, priors, targets)
loss = cfg['loc_weight'] * loss_l + loss_c + loss_landm
loss.backward()
optimizer.step()
load_t1 = time.time()
batch_time = load_t1 - load_t0
eta = int(batch_time * (max_iter - iteration))
print(
'Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || Loc: {:.4f} Cla: {:.4f} Landm: {:.4f} || LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}'
.format(epoch, max_epoch, (iteration % epoch_size) + 1, epoch_size,
iteration + 1, max_iter, loss_l.item(), loss_c.item(),
loss_landm.item(), lr, batch_time,
str(datetime.timedelta(seconds=eta))))
del images, targets
import gc
gc.collect()
torch.cuda.empty_cache()
ITER_SAVE_UNIT = 50
if iteration % ITER_SAVE_UNIT == 0:
# Save a loss
with open('./results/loss.txt', 'a') as f:
f.write(f'{loss.item():.4f}\n')
# Plot a graph
if iteration % (ITER_SAVE_UNIT * 10) == 0:
losses.append(loss.item())
iters = [iter * ITER_SAVE_UNIT for iter in range(len(losses))]
plt.plot(iters, losses)
plt.xlabel('Iteration')
plt.ylabel('Loss')
plt.title('Training Losses')
plt.savefig('./results/loss.png')
torch.save(net.state_dict(), save_folder + cfg['name'] + '_Final.pth')
