def train(train_loader, model, criterion, optimizer, epoch, epochs, device, logger, log_freq):
batch_time = utils.AverageMeter()
data_time = utils.AverageMeter()
losses = utils.AverageMeter()
acc_mean = utils.AverageMeter()
total_iter = len(train_loader) * train_loader.batch_size * epoch
# switch to train mode
model.train()
end = time.time()
pbar = tqdm(train_loader, unit='batches')
for i, (input, target) in enumerate(pbar):
# measure data loading time
data_time.update(time.time() - end)
input = input.to(device)
target = target.to(device)
# compute output
output = model(input)
if len(output.shape) > 2:
output = output.view(output.size(0), -1)
loss = criterion(output, target)
# measure accuracy and record loss
acc_ = get_acc(target, output)
losses.update(loss.item(), input.size(0))
acc_mean.update(acc_, input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
total_iter += train_loader.batch_size
# Batch logs
pbar.set_description(
'TRAINING: Epoch: {} / {}; '
'Timing: [Data: {batch_time.val:.3f} ({batch_time.avg:.3f}), '
'Batch: {data_time.val:.3f} ({data_time.avg:.3f})]; '
'Loss {loss.val:.4f} ({loss.avg:.4f}); '
'acc_mean {acc_mean.val:.3f} ({acc_mean.avg:.3f}); '.format(
epoch + 1, epochs, batch_time=batch_time, data_time=data_time,
loss=losses, acc_mean=acc_mean))
if logger and i % log_freq == 0:
logger.add_scalars('batch', {'loss': losses.val}, total_iter)
# logger.add_scalars('batch/acc', {'acc_mean': acc_mean.val}, total_iter)
# Epoch logs
if logger:
logger.add_scalars('epoch/acc', {'train': acc_mean.avg}, epoch)
logger.add_scalars('epoch/loss', {'train': losses.avg}, epoch)
def train(train_loader, model, criterion, optimizer, epoch, device, logger,
log_freq):
batch_time = utils.AverageMeter()
data_time = utils.AverageMeter()
losses = utils.AverageMeter()
total_iter = len(train_loader) * train_loader.batch_size * epoch
iters = len(train_loader) * train_loader.batch_size * (epoch + 1)
# switch to train mode
model.train()
# set time
end = time.time()
for i, (inputs, targets) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
batch_size = inputs.size(0)
inputs = inputs.to(device)
targets = targets.to(device)
# compute output
optimizer.zero_grad()
output = model(inputs)
# compute gradient and do OPTIMIZER step
loss_sum = 0
for n in range(len(output)):
loss_sum += criterion(input=output[n], target=targets)
loss_sum.backward()
optimizer.step()
# update metrics
losses.update(loss_sum.item())
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
total_iter += train_loader.batch_size
# Batch logs (reworked to semseg style)
if logger and i % log_freq == 0:
fmt_str = "Iter [{:d}/{:d}] Loss: {:.4f} Time/Image: {:.4f}"
print_str = fmt_str.format(total_iter, iters, losses.avg,
batch_time.avg / batch_size)
print(print_str)
logger.add_scalars('iterations', {'loss': losses.val}, total_iter)
# Epoch logs
if logger:
logger.add_scalars('epoch_loss', {'train_loss': losses.avg}, epoch)
def validate(val_loader, model, criterion, epoch, epochs, device, logger, running_metrics_val):
batch_time = utils.AverageMeter()
val_losses = utils.AverageMeter()
# switch to evaluate mode
model.eval()
with torch.no_grad():
end = time.time()
pbar = tqdm(val_loader, unit='batches')
for i, (input, target) in enumerate(pbar):
input = input.to(device)
target = target.to(device)
# compute output
output = model(input)
if len(output.shape) > 2:
output = output.view(output.size(0), -1)
val_loss = criterion(target, output)
# measure accuracy and record loss
val_losses.update(val_loss.item(), input.size(0))
trues = target.data.cpu()
preds = output.data.cpu()
running_metrics_val.update(trues, preds)
acc_mean = np.mean(running_metrics_val.acc_cls)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
pbar.set_description(
'VALIDATION: Epoch: {} / {}; '
'Timing: [Batch: {batch_time.val:.3f} ({batch_time.avg:.3f})]; '
'Loss {loss.val:.3f} ({loss.avg:.3f}); '
'acc_mean {acc_mean:.3f}; '.format(
epoch + 1, epochs, batch_time=batch_time,
loss=val_losses, acc_mean=acc_mean))
# Epoch logs
if logger:
class_dic = val_loader.dataset.class_dic
logger.add_scalars('epoch/acc', {'val': acc_mean}, epoch)
class_acc = running_metrics_val.acc_cls
print("val_loss: {}, mean_acc: {}".format(val_losses.avg, acc_mean))
for k, v in enumerate(class_acc):
logger.add_scalar('val_metrics/{}'.format(class_dic[k]), v, epoch)
logger.add_scalars('epoch/loss', {'val': val_losses.avg}, epoch)
return val_losses.avg, acc_mean
def validate(val_loader, model, criterion, epoch, epochs, device, logger, log_freq):
batch_time = utils.AverageMeter()
losses = utils.AverageMeter()
top1 = utils.AverageMeter()
top5 = utils.AverageMeter()
# switch to evaluate mode
model.eval()
with torch.no_grad():
end = time.time()
pbar = tqdm(val_loader, unit='batches')
for i, (input, target) in enumerate(pbar):
input = input.to(device)
target = target.to(device)
# compute output
output = model(input)
if len(output.shape) > 2:
output = output.view(output.size(0), -1)
loss = criterion(output, target)
# measure accuracy and record loss
prec1, prec5 = accuracy(output, target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(prec1[0], input.size(0))
top5.update(prec5[0], input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
pbar.set_description(
'VALIDATION: Epoch: {} / {}; '
'Timing: [Batch: {batch_time.val:.3f} ({batch_time.avg:.3f})]; '
'Loss {loss.val:.3f} ({loss.avg:.3f}); '
'Prec@1 {top1.val:.3f} ({top1.avg:.3f}); '
'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
epoch + 1, epochs, batch_time=batch_time,
loss=losses, top1=top1, top5=top5))
# print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
# .format(top1=top1, top5=top5))
# Epoch logs
if logger:
logger.add_scalars('epoch/acc/val', {'top1': top1.avg, 'top5': top5.avg}, epoch)
return losses.avg, top1.avg
def validate(val_loader, model, criterion, epoch, epochs, device, logger,
running_metrics_val):
batch_time = utils.AverageMeter()
val_los = utils.AverageMeter()
# switch to evaluate mode
model.eval()
with torch.no_grad():
end = time.time()
pbar = tqdm(val_loader, unit='batches')
for i, (inputs, targets, images) in enumerate(pbar):
inputs = inputs.to(device)
targets = targets.to(device)
images = images.to(device)
# compute output and loss
output = model(inputs)
loss_sum = 0
for n in range(len(output)):
loss_sum += criterion(input=output[n], target=targets)
# update metrics
pred = torch.zeros_like(output[0])
for n in range(len(output)):
pred += output[n].data
pred = pred / len(output)
pred = pred.max(1)[1]
gt = targets.data
running_metrics_val.update(gt, pred)
val_los.update(loss_sum.item())
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
pbar.set_description(
'VALIDATION: Epoch: {} / {}; '
'Timing: [Batch: {batch_time.val:.3f} ({batch_time.avg:.3f})]; '
'Loss {loss.val:.3f} ({loss.avg:.3f}); '.format(
epoch + 1, epochs, batch_time=batch_time, loss=val_los))
# Metrics
score, class_iou = running_metrics_val.get_scores()
running_metrics_val.reset()
# Epoch logs
if logger:
report = []
for k, v in score.items():
k, v = k, "{:.3f}".format(v.item())
report.append((k, v))
logger.add_scalar('val_metrics/{}'.format(k.strip()), float(v),
epoch)
print(k, v)
print("val_loss: {:.3f}".format(val_los.avg))
report.append("val_loss: {:.3f}".format(val_los.avg))
for k, v in class_iou.items():
k, v = k, "{:.3f}".format(v.item())
report.append((k, v))
logger.add_scalar('class_metrics/cls_{}'.format(k), float(v),
epoch)
logger.add_scalars('epoch_loss', {'val_loss': val_los.avg}, epoch)
# Log images
blend = seg_utils.blend_seg_label(images, pred)
gt_blend = seg_utils.blend_seg_label(images, targets)
pred = pred.view(pred.shape[0], 1, pred.shape[1],
pred.shape[2]).repeat(1, 3, 1, 1).float()
grid = seg_utils.make_grid(images, gt_blend, pred, blend,
min(8, val_loader.batch_size))
logger.add_image('%d/rec' % (256), grid, epoch)
return val_los.avg, score, report
def validate(val_loader, model, device, val_dir, batch_size):
batch_time = utils.AverageMeter()
# switch to evaluate mode
model.eval()
with torch.no_grad():
end = time.time()
pbar = tqdm(val_loader, unit='batches')
c = 0
new_list = []
for j, (inputs, im_names, hs, ws) in enumerate(pbar):
_, _, h, w = inputs.size()
val_inputs = inputs.to(device)
# compute output and loss
output = model(val_inputs)
# getting images and predictions as numpy arrays
pred = output.data.max(1)[1].cpu().numpy()
hs = hs.cpu().numpy()
ws = ws.cpu().numpy()
for i in range(pred.shape[0]):
mask = Image.fromarray(pred[i].astype(np.uint8))
imP = mask.convert('RGB').convert('P',
palette=Image.ADAPTIVE,
colors=3)
reverse_colors = np.array(imP)
# we want to exclude poor quality predictions, practically found threshold == 7000
a = np.sum(reverse_colors[reverse_colors == 1]) # face class
b = np.count_nonzero(reverse_colors == 0) # hair class
if len(np.unique(reverse_colors)) == 3 and a + b > (0.15 * h *
w):
reverse_colors[reverse_colors == 0] = 3
reverse_colors[reverse_colors == 2] = 0
reverse_colors[reverse_colors == 3] = 2
imP = Image.fromarray(reverse_colors, mode='P')
imP.putpalette([
0,
0,
0, # index 0 is black (background)
0,
255,
0, # index 1 is green (face)
255,
0,
0, # index 2 is red (hair)
])
# resize to match original image
h, w = hs[i], ws[i]
imP.resize(size=(h, w))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
folder = os.path.join(val_dir, 'masks',
im_names[i].split('/')[0])
if not os.path.exists(folder):
os.makedirs(folder)
path_save = os.path.join(
folder, "{}.png".format(im_names[i].split('/')[-1]))
# save data we need
imP.save(path_save)
new_list.append(im_names[i])
c += 1
pbar.set_description(
'VALIDATION: Iter: {} / {}; '
'Timing: [Batch: {batch_time.val:.3f} ({batch_time.avg:.3f})]; '
.format(c,
len(val_loader) * batch_size,
batch_time=batch_time))
path_t = os.path.join(val_dir, 'masks.txt')
with open(path_t, 'w') as t:
for item in new_list:
t.write("%s\n" % item)