<reponame>bchiu/Simple-CenterNet<gh_stars>1-10
from models import centernet
from utils import common
from data import dataset

import torch
import torch.cuda.amp
from torch.utils.tensorboard import SummaryWriter

import argparse
import os

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--step-batch-size', default=32, type=int, help='Batch size for step(optimization)')
    parser.add_argument('--forward-batch-size', default=32, type=int, help='Batch size for forward')
    
    parser.add_argument('--img-w', default=512, type=int)
    parser.add_argument('--img-h', default=512, type=int)

    parser.add_argument('--weights', type=str, default="", help='load weights to resume training')
    parser.add_argument('--total-epoch', type=int, default=70, help='total_epoch')

    parser.add_argument('--data', type=str, default="./data/voc0712.yaml")
    parser.add_argument('--num-workers', default=8, type=int, help='Number of workers used in dataloading')
    parser.add_argument('--save-folder', default='./weights', type=str, help='where you save weights')
    parser.add_argument('--seed', default=7777, type=int)

    opt = parser.parse_args()
    common.setup_seed(opt.seed)
    common.mkdir(dir=opt.save_folder, remove_existing_dir=False)
    
    dataset_dict = common.parse_yaml(opt.data)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    model = centernet.CenterNet(num_classes=len(dataset_dict['classes']), pretrained_backbone=True)
    model = model.to(device=device)
    
    training_set = dataset.DetectionDataset(root=dataset_dict['root'], 
                                            dataset_name=dataset_dict['dataset_name'],
                                            set="train",
                                            num_classes=len(dataset_dict['classes']),
                                            img_w=opt.img_w, img_h=opt.img_h,
                                            use_augmentation=True,
                                            keep_ratio=True)

    training_set_loader = torch.utils.data.DataLoader(training_set, 
                                                      opt.forward_batch_size,
                                                      num_workers=opt.num_workers,
                                                      shuffle=True,
                                                      collate_fn=dataset.collate_fn,
                                                      pin_memory=True,
                                                      drop_last=True)

    iters_to_accumulate = max(round(opt.step_batch_size/opt.forward_batch_size), 1)
    initial_lr = 5e-4 * (opt.step_batch_size/128)
    optimizer = torch.optim.Adam(model.parameters(), lr=initial_lr)
    for param_group in optimizer.param_groups:
        param_group['lr'] = 0.
    
    iterations_per_epoch = len(training_set_loader)
    total_iteration =  iterations_per_epoch * opt.total_epoch

    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_iteration)
    warmup_iteration = 1000
    scaler = torch.cuda.amp.GradScaler()
    
    start_epoch = 0
    if os.path.isfile(opt.weights):
        checkpoint = torch.load(opt.weights)
        start_epoch = checkpoint['epoch'] + 1
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        scaler.load_state_dict(checkpoint['scaler_state_dict'])
    
    writer = SummaryWriter()
    for epoch in range(start_epoch, opt.total_epoch):
        model.train()
        for i, batch_data in enumerate(training_set_loader):
            
            n_iteration = (iterations_per_epoch * epoch) + i
            
            batch_img = batch_data["img"].to(device)
            batch_label = batch_data["label"]
            
            #forward
            with torch.cuda.amp.autocast():
                batch_output = model(batch_img)
                loss, losses = model.compute_loss(batch_output, batch_label)
            
            writer.add_scalar('train/loss_offset_xy', losses[0].item(), n_iteration)
            writer.add_scalar('train/loss_wh', losses[1].item(), n_iteration)
            writer.add_scalar('train/loss_class_heatmap', losses[2].item(), n_iteration)
            writer.add_scalar('train/loss', loss.item(), n_iteration)
            writer.add_scalar('train/lr', common.get_lr(optimizer), n_iteration)

            #backword
            loss = loss / iters_to_accumulate
            scaler.scale(loss).backward()

            if (n_iteration + 1) % iters_to_accumulate == 0:
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad()

            if n_iteration > warmup_iteration:
                scheduler.step()
            else:
                lr = initial_lr * float(n_iteration) / warmup_iteration
                for param_group in optimizer.param_groups:
                    param_group['lr'] = lr
        
        checkpoint = {
        'epoch': epoch,# zero indexing
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict' : optimizer.state_dict(),
        'scheduler_state_dict' : scheduler.state_dict(),
        'scaler_state_dict' : scaler.state_dict(),
        'mAP': 0.,
        'best_mAP': 0.
        }
        
        torch.save(checkpoint, os.path.join(opt.save_folder, 'epoch_' + str(epoch + 1) + '.pth'))
        

from models import centernet
from utils import common
from data import dataset

import torch
import torch.cuda.amp
from torch.utils.tensorboard import SummaryWriter

import argparse
import os

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--step-batch-size', default=32, type=int, help='Batch size for step(optimization)')
    parser.add_argument('--forward-batch-size', default=32, type=int, help='Batch size for forward')
    
    parser.add_argument('--img-w', default=512, type=int)
    parser.add_argument('--img-h', default=512, type=int)

    parser.add_argument('--weights', type=str, default="", help='load weights to resume training')
    parser.add_argument('--total-epoch', type=int, default=70, help='total_epoch')

    parser.add_argument('--data', type=str, default="./data/voc0712.yaml")
    parser.add_argument('--num-workers', default=8, type=int, help='Number of workers used in dataloading')
    parser.add_argument('--save-folder', default='./weights', type=str, help='where you save weights')
    parser.add_argument('--seed', default=7777, type=int)

    opt = parser.parse_args()
    common.setup_seed(opt.seed)
    common.mkdir(dir=opt.save_folder, remove_existing_dir=False)
    
    dataset_dict = common.parse_yaml(opt.data)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    model = centernet.CenterNet(num_classes=len(dataset_dict['classes']), pretrained_backbone=True)
    model = model.to(device=device)
    
    training_set = dataset.DetectionDataset(root=dataset_dict['root'], 
                                            dataset_name=dataset_dict['dataset_name'],
                                            set="train",
                                            num_classes=len(dataset_dict['classes']),
                                            img_w=opt.img_w, img_h=opt.img_h,
                                            use_augmentation=True,
                                            keep_ratio=True)

    training_set_loader = torch.utils.data.DataLoader(training_set, 
                                                      opt.forward_batch_size,
                                                      num_workers=opt.num_workers,
                                                      shuffle=True,
                                                      collate_fn=dataset.collate_fn,
                                                      pin_memory=True,
                                                      drop_last=True)

    iters_to_accumulate = max(round(opt.step_batch_size/opt.forward_batch_size), 1)
    initial_lr = 5e-4 * (opt.step_batch_size/128)
    optimizer = torch.optim.Adam(model.parameters(), lr=initial_lr)
    for param_group in optimizer.param_groups:
        param_group['lr'] = 0.
    
    iterations_per_epoch = len(training_set_loader)
    total_iteration =  iterations_per_epoch * opt.total_epoch

    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_iteration)
    warmup_iteration = 1000
    scaler = torch.cuda.amp.GradScaler()
    
    start_epoch = 0
    if os.path.isfile(opt.weights):
        checkpoint = torch.load(opt.weights)
        start_epoch = checkpoint['epoch'] + 1
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        scaler.load_state_dict(checkpoint['scaler_state_dict'])
    
    writer = SummaryWriter()
    for epoch in range(start_epoch, opt.total_epoch):
        model.train()
        for i, batch_data in enumerate(training_set_loader):
            
            n_iteration = (iterations_per_epoch * epoch) + i
            
            batch_img = batch_data["img"].to(device)
            batch_label = batch_data["label"]
            
            #forward
            with torch.cuda.amp.autocast():
                batch_output = model(batch_img)
                loss, losses = model.compute_loss(batch_output, batch_label)
            
            writer.add_scalar('train/loss_offset_xy', losses[0].item(), n_iteration)
            writer.add_scalar('train/loss_wh', losses[1].item(), n_iteration)
            writer.add_scalar('train/loss_class_heatmap', losses[2].item(), n_iteration)
            writer.add_scalar('train/loss', loss.item(), n_iteration)
            writer.add_scalar('train/lr', common.get_lr(optimizer), n_iteration)

            #backword
            loss = loss / iters_to_accumulate
            scaler.scale(loss).backward()

            if (n_iteration + 1) % iters_to_accumulate == 0:
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad()

            if n_iteration > warmup_iteration:
                scheduler.step()
            else:
                lr = initial_lr * float(n_iteration) / warmup_iteration
                for param_group in optimizer.param_groups:
                    param_group['lr'] = lr
        
        checkpoint = {
        'epoch': epoch,# zero indexing
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict' : optimizer.state_dict(),
        'scheduler_state_dict' : scheduler.state_dict(),
        'scaler_state_dict' : scaler.state_dict(),
        'mAP': 0.,
        'best_mAP': 0.
        }
        
        torch.save(checkpoint, os.path.join(opt.save_folder, 'epoch_' + str(epoch + 1) + '.pth'))
#forward #backword # zero indexing