class Epoch:
def __init__(self,
model,
num_classes,
losses,
stage_name,
device='cpu',
verbose=True,
EvaluatorIn=None):
self.model = model
self.losses = losses
self.stage_name = stage_name
self.verbose = verbose
self.device = device
self._to_device()
if EvaluatorIn is None:
self.evaluator = Evaluator(num_classes, device)
else:
self.evaluator = EvaluatorIn
def _to_device(self):
for loss in self.losses:
loss.to(self.device)
self.model.to(self.device)
@classmethod
def _format_logs(self, logs):
str_logs = ['{} - {:.4}'.format(k, v) for k, v in logs.items()]
s = ', '.join(str_logs)
return s
def batch_update(self, x, y):
raise NotImplementedError
def on_epoch_start(self):
pass
def run(self, dataloader):
self.on_epoch_start()
logs = {}
loss_meter = AverageValueMeter()
losses_meters = {
loss.__name__: AverageValueMeter()
for loss in self.losses
}
with tqdm(dataloader,
desc=self.stage_name,
file=sys.stdout,
disable=not (self.verbose)) as iterator:
for sample in iterator:
x = sample['image']
y = sample['label']
if 'image_class' in sample:
image_class = sample['image_class']
else:
image_class = None
x, y = x.to(self.device), y.to(self.device)
loss, losses, y_pred = self.batch_update(
x, y, image_class, sample)
# update losses logs
loss_value = loss.cpu().detach().numpy()
loss_meter.add(loss_value)
loss_logs = {'total_loss': loss_meter.mean}
logs.update(loss_logs)
for loss_fn in self.losses:
loss_value = loss_fn(y_pred, y, x, image_class)
if type(loss_value) == torch.Tensor:
loss_value = loss_value.cpu().detach().numpy()
losses_meters[loss_fn.__name__].add(loss_value)
losses_logs = {k: v.mean for k, v in losses_meters.items()}
logs.update(losses_logs)
if self.verbose:
s = self._format_logs(logs)
iterator.set_postfix_str(s)
logs['acc'] = self.evaluator.Pixel_Accuracy()
logs['fscore'] = self.evaluator.f_score()
logs['jaccard'] = self.evaluator.Jaccard()
logs['miou'] = self.evaluator.Mean_Intersection_over_Union()
logs['acc_class'] = self.evaluator.Pixel_Accuracy_Class()
return logs
class Trainer(object):
def __init__(self, args):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(args.logdir)
self.writer = self.summary.create_summary()
# Define Dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
dltrain = DLDataset('trainval', "./data/pascal_voc_seg/tfrecord/")
dlval = DLDataset('val', "./data/pascal_voc_seg/tfrecord/")
# dltrain = DLDataset('trainval', "./data/pascal_voc_seg/VOCdevkit/VOC2012/")
# dlval = DLDataset('val', "./data/pascal_voc_seg/VOCdevkit/VOC2012/")
self.train_loader = DataLoader(dltrain,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
self.val_loader = DataLoader(dlval,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
# Define network
model = Deeplab()
train_params = [{
'params': model.get_1x_lr_params(),
'lr': args.lr
}, {
'params': model.get_10x_lr_params(),
'lr': args.lr * 10
}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=args.nesterov)
# Define Criterion
# whether to use class balanced weights
self.criterion = nn.CrossEntropyLoss(ignore_index=255).cuda()
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(21)
# Define lr scheduler
self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer=optimizer)
# Using cuda
# if args.cuda:
# self.model = torch.nn.DataParallel(self.model)
self.model = self.model.cuda()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# Clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
def training(self, epoch):
train_loss = 0.0
self.model.train()
tbar = tqdm(self.train_loader)
num_img_tr = len(self.train_loader)
for i, (image, target) in enumerate(tbar):
if self.args.cuda:
image, target = image.cuda(), target.cuda()
self.optimizer.zero_grad()
output = self.model(image)
loss = self.criterion(output, target.long())
loss.backward()
self.optimizer.step()
train_loss += loss.item()
tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
self.writer.add_scalar('train/total_loss_iter', loss.item(),
i + num_img_tr * epoch)
# Show 10 * 3 inference results each epoch
# if i % (num_img_tr // 10) == 0:
if i % 10 == 0:
global_step = i + num_img_tr * epoch
self.summary.visualize_image(self.writer, self.args.dataset,
image, target, output,
global_step)
self.scheduler.step(train_loss)
self.writer.add_scalar('train/total_loss_epoch', train_loss, epoch)
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.args.batch_size + image.data.shape[0]))
print('Loss: %.3f' % train_loss)
if self.args.no_val:
# save checkpoint every epoch
is_best = False
self.saver.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': self.model.module.state_dict(),
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
}, is_best)
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
tbar = tqdm(self.val_loader, desc='\r')
test_loss = 0.0
for i, (image, target) in enumerate(tbar):
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
loss = self.criterion(output, target.long())
test_loss += loss.item()
tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# Fast test during the training
Acc = self.evaluator.Pixel_Accuracy()
Acc_class = self.evaluator.Pixel_Accuracy_Class()
mIoU = self.evaluator.Mean_Intersection_over_Union()
FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch)
self.writer.add_scalar('val/mIoU', mIoU, epoch)
self.writer.add_scalar('val/Acc', Acc, epoch)
self.writer.add_scalar('val/Acc_class', Acc_class, epoch)
self.writer.add_scalar('val/fwIoU', FWIoU, epoch)
print('Validation:')
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.args.batch_size + image.data.shape[0]))
print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(
Acc, Acc_class, mIoU, FWIoU))
print('Loss: %.3f' % test_loss)
new_pred = mIoU
if new_pred > self.best_pred:
is_best = True
self.best_pred = new_pred
self.saver.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': self.model.state_dict(),
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
}, is_best)
