train.py

import datetime
import os
import traceback

import numpy as np
import torch
import yaml
from torch import nn
from torch.utils.data import DataLoader
from torchvision import transforms
from tqdm.autonotebook import tqdm

from backbone import EfficientDetBackbone
from efficientdet.dataset import Resizer, Normalizer, collater
from efficientdet.loss import FocalLoss
from utils.sync_batchnorm import patch_replication_callback
from utils.utils import replace_w_sync_bn, CustomDataParallel, get_last_weights, init_weights
from dataprocessor import DatawAnnotation, NoiseAdder


os.environ["CUDA_VISIBLE_DEVICES"] = '0'
project = 'numbers'


class Params:
    def __init__(self, project_file):
        self.params = yaml.safe_load(open(project_file).read())

    def __getattr__(self, item):
        return self.params.get(item, None)


class ModelWithLoss(nn.Module):
    def __init__(self, model, debug=False):
        super().__init__()
        self.criterion = FocalLoss()
        self.model = model
        self.debug = debug

    def forward(self, imgs, annotations, obj_list=None):
        _, regression, classification, anchors = self.model(imgs)
        if self.debug:
            cls_loss, reg_loss = self.criterion(classification, regression,
                                                anchors, annotations,
                                                imgs=imgs, obj_list=obj_list)
        else:
            cls_loss, reg_loss = self.criterion(classification, regression,
                                                anchors, annotations)
        return cls_loss, reg_loss


def train(project):
    params = Params(f'{project}.yml')

    if params.num_gpus == 0:
        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

    if torch.cuda.is_available():
        torch.cuda.manual_seed(42)
    else:
        torch.manual_seed(42)

    params.saved_path = params.saved_path + f'/{params.project_name}/'
    params.log_path = params.log_path + f'/{params.project_name}/tensorboard/'
    os.makedirs(params.log_path, exist_ok=True)
    os.makedirs(params.saved_path, exist_ok=True)

    training_params = {'batch_size': params.batch_size,
                       'shuffle': True,
                       'drop_last': True,
                       'collate_fn': collater,
                       'num_workers': params.num_workers}

    val_params = {'batch_size': params.batch_size,
                  'shuffle': False,
                  'drop_last': True,
                  'collate_fn': collater,
                  'num_workers': params.num_workers}

    input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536, 1536]
    if params.custom_size:
        img_size = params.custom_size
    else:
        img_size = input_sizes[params.compound_coef]

    training_set = DatawAnnotation(root_dir='train/',
                                   anno_dir='annotation_train.csv',
                                   transform=transforms.Compose([
                                       Normalizer(mean=params.mean, std=params.std),
                                       NoiseAdder(),
                                       Resizer(img_size)
                                   ]))
    training_generator = DataLoader(training_set, **training_params)

    val_set = DatawAnnotation(root_dir='val/',
                              anno_dir='annotation_val.csv',
                              transform=transforms.Compose([
                                  Normalizer(mean=params.mean, std=params.std),
                                  Resizer(img_size)
                              ]))
    val_generator = DataLoader(val_set, **val_params)

    model = EfficientDetBackbone(num_classes=len(params.obj_list),
                                 compound_coef=params.compound_coef,
                                 ratios=eval(params.anchors_ratios),
                                 scales=eval(params.anchors_scales))

    # load last weights
    if params.load_weights is not None:
        if params.load_weights.endswith('.pth'):
            weights_path = params.load_weights
        else:
            weights_path = get_last_weights(params.saved_path)
        try:
            last_step = int(os.path.basename(weights_path).split('_')[-1].split('.')[0])
        except:
            last_step = 0

        try:
            ret = model.load_state_dict(torch.load(weights_path), strict=False)
        except RuntimeError as e:
            print(f'[Warning] Ignoring {e}')
            print(
                '[Warning] Don\'t panic if you see this, this might be because '
                + 'you load a pretrained weights with different number of classes. '
                + 'The rest of the weights should be loaded already.')

        print(f'[Info] loaded weights: {os.path.basename(weights_path)},'
              + 'resuming checkpoint from step: {last_step}')
    else:
        last_step = 0
        print('[Info] initializing weights...')
        init_weights(model)

    # freeze backbone if train head_only
    if params.head_only:
        def freeze_backbone(m):
            classname = m.__class__.__name__
            for ntl in ['EfficientNet', 'BiFPN']:
                if ntl in classname:
                    for param in m.parameters():
                        param.requires_grad = False

        model.apply(freeze_backbone)
        print('[Info] freezed backbone')

    # https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
    # apply sync_bn when using multiple gpu and batch_size per gpu
    #  is lower than 4, useful when gpu memory is limited.
    # because when bn is disable, the training will be very unstable
    #  or slow to converge, apply sync_bn can solve it,
    # by packing all mini-batch across all gpus as one batch and normalize,
    #  then send it back to all gpus.
    # but it would also slow down the training by a little bit.
    if params.num_gpus > 1 and params.batch_size // params.num_gpus < 4:
        model.apply(replace_w_sync_bn)
        use_sync_bn = True
    else:
        use_sync_bn = False

    # warp the model with loss function, to reduce the memory usage and speedup
    model = ModelWithLoss(model, debug=params.debug)

    if params.num_gpus > 0:
        model = model.cuda()
        if params.num_gpus > 1:
            model = CustomDataParallel(model, params.num_gpus)
            if use_sync_bn:
                patch_replication_callback(model)

    if params.optim == 'adamw':
        optimizer = torch.optim.AdamW(model.parameters(), params.lr)
    else:
        optimizer = torch.optim.SGD(model.parameters(), params.lr, momentum=0.9,
                                    nesterov=True)

    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3,
                                                           verbose=True)

    epoch = 0
    best_loss = 1e5
    best_epoch = 0
    step = max(0, last_step)
    model.train()

    num_iter_per_epoch = len(training_generator)

    with open('log.txt', 'a+') as f:
        f.write(datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))

    try:
        for epoch in range(params.num_epochs):
            total_closs = 0
            total_rloss = 0
            last_epoch = step // num_iter_per_epoch
            if epoch < last_epoch:
                continue

            epoch_loss = []
            progress_bar = tqdm(training_generator)
            for iter, data in enumerate(progress_bar):
                if iter < step - last_epoch * num_iter_per_epoch:
                    progress_bar.update()
                    continue
                try:
                    imgs = data['img']
                    annot = data['annot']

                    if params.num_gpus == 1:
                        # if only one gpu, just send it to cuda:0
                        # elif multiple gpus, send it to multiple gpus in
                        # CustomDataParallel, not here
                        imgs = imgs.cuda()
                        annot = annot.cuda()

                    optimizer.zero_grad()
                    cls_loss, reg_loss = model(imgs, annot,
                                               obj_list=params.obj_list)
                    cls_loss = cls_loss.mean()
                    reg_loss = reg_loss.mean()

                    loss = cls_loss + reg_loss
                    total_closs += cls_loss
                    total_rloss += reg_loss
                    if loss == 0 or not torch.isfinite(loss):
                        continue

                    loss.backward()
                    # torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1)
                    optimizer.step()

                    epoch_loss.append(float(loss))

                    progress_bar.set_description(
                        'Step: {}. Epoch: {}/{}. Iteration: {}/{}. Cls loss: {:.5f}. Reg loss: {:.5f}. Total loss: {:.5f}'.format(
                            step, epoch, params.num_epochs, iter + 1,
                            num_iter_per_epoch, cls_loss.item(),
                            reg_loss.item(), loss.item()))

                    # log learning_rate
                    current_lr = optimizer.param_groups[0]['lr']

                    step += 1

                    if step % params.save_interval == 0 and step > 0:
                        save_checkpoint(model,
                                        f'efficientdet-d{params.compound_coef}_{epoch}_{step}.pth',
                                        params.saved_path)
                        print('checkpoint...')

                except Exception as e:
                    print('[Error]', traceback.format_exc())
                    print(e)
                    continue
            scheduler.step(np.mean(epoch_loss))

            with open('log.txt', 'a+') as f:
                f.write(
                    'Train. Epoch: {}/{}. Classification loss: {:1.5f}. Regression loss: {:1.5f}. Total loss: {:1.5f}'.format(
                        epoch, params.num_epochs, total_closs, total_rloss,
                        total_closs + total_rloss))

            # Validation
            if epoch % params.val_interval == 0:
                model.eval()
                loss_regression_ls = []
                loss_classification_ls = []
                for iter, data in enumerate(val_generator):
                    with torch.no_grad():
                        imgs = data['img']
                        annot = data['annot']

                        if params.num_gpus == 1:
                            imgs = imgs.cuda()
                            annot = annot.cuda()

                        cls_loss, reg_loss = model(imgs, annot,
                                                   obj_list=params.obj_list)
                        cls_loss = cls_loss.mean()
                        reg_loss = reg_loss.mean()

                        loss = cls_loss + reg_loss
                        if loss == 0 or not torch.isfinite(loss):
                            continue

                        loss_classification_ls.append(cls_loss.item())
                        loss_regression_ls.append(reg_loss.item())

                cls_loss = np.mean(loss_classification_ls)
                reg_loss = np.mean(loss_regression_ls)
                loss = cls_loss + reg_loss

                print(
                    'Val. Epoch: {}/{}. Classification loss: {:1.5f}. Regression loss: {:1.5f}. Total loss: {:1.5f}'.format(
                        epoch, params.num_epochs, cls_loss, reg_loss, loss))

                with open('log.txt', 'a+') as f:
                    f.write(
                        'Val. Epoch: {}/{}. Classification loss: {:1.5f}. Regression loss: {:1.5f}. Total loss: {:1.5f}'.format(
                            epoch, params.num_epochs, total_closs, total_rloss,
                            total_closs + total_rloss))

                if loss + params.es_min_delta < best_loss:
                    best_loss = loss
                    best_epoch = epoch

                    save_checkpoint(model,
                                    f'efficientdet-d{params.compound_coef}_{epoch}_{step}.pth',
                                    params.saved_path)

                model.train()

                # Early stopping
                if epoch - best_epoch > params.es_patience > 0:
                    print('[Info] Stop training at epoch {}. The lowest loss achieved is {}'.format(epoch, best_loss))
                    break
    except KeyboardInterrupt:
        save_checkpoint(model,
                        f'efficientdet-d{params.compound_coef}_{epoch}_{step}.pth',
                        params.saved_path)


def save_checkpoint(model, name, save_path):
    if isinstance(model, CustomDataParallel):
        torch.save(model.module.model.state_dict(), os.path.join(save_path, name))
    else:
        torch.save(model.model.state_dict(), os.path.join(save_path, name))


if __name__ == '__main__':
    train(project)