train_ggcnn_multi_gpus.py

"""
written by ryanreadbooks
date: 2021/10/23
"""

import datetime
import os
import sys
import argparse
import logging
import time

import cv2

import torch
import torch.utils.data
import torch.optim as optim
import torch.distributed

from torchsummary import summary

import tensorboardX

from utils.visualisation.gridshow import gridshow

from utils.dataset_processing import evaluation
from utils.data import get_dataset
from models import get_network
from models.common import post_process_output

logging.basicConfig(level=logging.INFO)


def logout(msg):
    if args.local_rank == 0:
        logging.info(msg)


def parse_args():
    parser = argparse.ArgumentParser(description='Train GG-CNN')

    # Network
    parser.add_argument('--network', type=str, default='ggcnn', help='Network Name in models')

    # Dataset & Data & Training
    parser.add_argument('--dataset', type=str, default='cornell', help='Dataset Name ("cornell" or "jaquard" or "graspnet1b")')
    parser.add_argument('--dataset-path', type=str, help='Path to dataset')
    parser.add_argument('--use-depth', type=int, default=1, help='Use Depth image for training (1/0)')
    parser.add_argument('--use-rgb', type=int, default=0, help='Use RGB image for training (0/1)')
    parser.add_argument('--split', type=float, default=0.9, help='Fraction of data for training (remainder is validation)')
    parser.add_argument('--ds-rotate', type=float, default=0.0,
                        help='Shift the start point of the dataset to use a different test/train split for cross validation.')
    parser.add_argument('--output-size', type=int, default=300,
                        help='the output size of the network, determining the cropped size of dataset images')

    parser.add_argument('--camera', type=str, default='realsense',
                        help='Which camera\'s data to use, only effective when using graspnet1b dataset')
    parser.add_argument('--scale', type=int, default=1,
                        help='the scale factor for the original images, only effective when using graspnet1b dataset')
    parser.add_argument('--ggcnn3backend', type=str, default='dla60up',
                        help='the backbone for GGCNN3, only effective when using GGCNN3 network')
    parser.add_argument('--num-workers', type=int, default=8, help='Dataset workers')
    parser.add_argument('--device', type=int, default=0, help='the gpu device number, starting from 0')

    parser.add_argument('--batch-size', type=int, default=8, help='Batch size')
    parser.add_argument('--epochs', type=int, default=50, help='Training epochs')
    parser.add_argument('--batches-per-epoch', type=int, default=1000, help='Batches per Epoch')
    parser.add_argument('--val-batches', type=int, default=250, help='Validation Batches')

    # Logging etc.
    parser.add_argument('--description', type=str, default='', help='Training description')
    parser.add_argument('--outdir', type=str, default='output/models/', help='Training Output Directory')
    parser.add_argument('--logdir', type=str, default='tensorboard/', help='Log directory')
    parser.add_argument('--vis', action='store_true', help='Visualise the training process')

    # multi gpus training settings
    parser.add_argument('--devices', type=str, default='0',
                        help="the gpu device number, starting from 0, separate with ','")
    parser.add_argument('--local_rank', type=int, default=-1, help='do not explicitly specify this parameter')
    args = parser.parse_args()
    return args


def validate(net, device, val_data, batches_per_epoch):
    """
    Run validation.
    :param net: Network
    :param device: Torch device
    :param val_data: Validation Dataset
    :param batches_per_epoch: Number of batches to run
    :return: Successes, Failures and Losses
    """
    net.eval()

    results = {
        'correct': 0,
        'failed': 0,
        'loss': 0,
        'losses': {

        }
    }

    ld = len(val_data)

    with torch.no_grad():
        batch_idx = 0
        while batch_idx < batches_per_epoch:
            for x, y, didx, rot, zoom_factor in val_data:
                batch_idx += 1
                if batches_per_epoch is not None and batch_idx >= batches_per_epoch:
                    break

                xc = x.to(device)
                yc = [yy.to(device) for yy in y]
                lossd = net.compute_loss(xc, yc)

                loss = lossd['loss']

                results['loss'] += loss.item() / ld
                for ln, l in lossd['losses'].items():
                    if ln not in results['losses']:
                        results['losses'][ln] = 0
                    results['losses'][ln] += l.item() / ld

                q_out, ang_out, w_out = post_process_output(lossd['pred']['pos'], lossd['pred']['cos'],
                                                            lossd['pred']['sin'], lossd['pred']['width'])

                s = evaluation.calculate_iou_match(q_out, ang_out,
                                                   val_data.dataset.get_gtbb(didx, rot, zoom_factor),
                                                   no_grasps=1,
                                                   grasp_width=w_out,
                                                   )

                if s:
                    results['correct'] += 1
                else:
                    results['failed'] += 1

    return results


def train(epoch, net, device, train_data, optimizer, batches_per_epoch, vis=False):
    """
    Run one training epoch
    :param epoch: Current epoch
    :param net: Network
    :param device: Torch device
    :param train_data: Training Dataset
    :param optimizer: Optimizer
    :param batches_per_epoch:  Data batches to train on
    :param vis:  Visualise training progress
    :return:  Average Losses for Epoch
    """
    results = {
        'loss': 0,
        'losses': {
        }
    }

    net.train()

    batch_idx = 0
    # Use batches per epoch to make training on different sized datasets (cornell/jacquard) more equivalent.
    while batch_idx < batches_per_epoch:
        for x, y, _, _, _ in train_data:
            batch_idx += 1
            if batch_idx >= batches_per_epoch:
                break

            xc = x.to(device)
            yc = [yy.to(device) for yy in y]
            lossd = net.compute_loss(xc, yc)

            loss = lossd['loss']

            if batch_idx % 100 == 0:
                logout('Epoch: {}, Batch: {}, Loss: {:0.4f}'.format(epoch, batch_idx, loss.item()))

            results['loss'] += loss.item()
            for ln, l in lossd['losses'].items():
                if ln not in results['losses']:
                    results['losses'][ln] = 0
                results['losses'][ln] += l.item()

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # Display the images
            if vis:
                imgs = []
                n_img = min(4, x.shape[0])
                for idx in range(n_img):
                    imgs.extend([x[idx,].numpy().squeeze()] + [yi[idx,].numpy().squeeze() for yi in y] + [
                        x[idx,].numpy().squeeze()] + [pc[idx,].detach().cpu().numpy().squeeze() for pc in lossd['pred'].values()])
                gridshow('Display', imgs,
                         [(xc.min().item(), xc.max().item()), (0.0, 1.0), (0.0, 1.0), (-1.0, 1.0), (0.0, 1.0)] * 2 * n_img,
                         [cv2.COLORMAP_BONE] * 10 * n_img, 10)
                cv2.waitKey(2)

    results['loss'] /= batch_idx
    for l in results['losses']:
        results['losses'][l] /= batch_idx

    return results


def run():
    # Vis window
    if args.vis:
        cv2.namedWindow('Display', cv2.WINDOW_NORMAL)

    if args.local_rank == 0:
        # Set-up output directories
        dt = datetime.datetime.now().strftime('%y%m%d_%H%M')
        net_desc = '{}_{}'.format(dt, '_'.join(args.description.split()))

        save_folder = os.path.join(args.outdir, net_desc)
        if not os.path.exists(save_folder):
            os.makedirs(save_folder)
        tb = tensorboardX.SummaryWriter(os.path.join(args.logdir, net_desc))

    # Load Dataset
    logout('Loading {} Dataset...'.format(args.dataset.title()))

    Dataset = get_dataset(args.dataset)
    train_dataset = Dataset(args.dataset_path, start=0.0, end=args.split, ds_rotate=args.ds_rotate,
                            output_size=args.output_size,
                            random_rotate=True, random_zoom=True,
                            include_depth=args.use_depth,
                            include_rgb=args.use_rgb,
                            camera=args.camera,
                            scale=args.scale)

    train_sampler = torch.utils.data.DistributedSampler(train_dataset)

    train_data = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        pin_memory=True,
        sampler=train_sampler
    )

    if args.local_rank == 0:
        val_dataset = Dataset(args.dataset_path, start=args.split, end=1.0, ds_rotate=args.ds_rotate,
                              output_size=args.output_size,
                              random_rotate=True, random_zoom=True,
                              include_depth=args.use_depth,
                              include_rgb=args.use_rgb,
                              camera=args.camera,
                              scale=args.scale)
        val_data = torch.utils.data.DataLoader(
            val_dataset,
            batch_size=1,  # do not modify
            shuffle=False,
            num_workers=args.num_workers // 2,
            pin_memory=True
        )

        logout(f'len of train loader = {len(train_data)}, len of eval loader = {len(val_data)}')
    logout('Done')

    # Load the network
    logout('Loading Network...')
    input_channels = 1 * args.use_depth + 3 * args.use_rgb
    ggcnn = get_network(args.network)

    if args.network == 'ggcnn3':
        net = ggcnn(input_channels=input_channels, backend=args.ggcnn3backend)
    else:
        net = ggcnn(input_channels=input_channels)

    devices = list(map(int, args.devices.split(',')))
    device = torch.device("cuda:{}".format(devices[args.local_rank]))

    torch.distributed.init_process_group(backend='nccl')
    torch.cuda.set_device(device=device)
    net = net.to(device)
    net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[device], output_device=device)
    net = net.module
    optimizer = optim.Adam(net.parameters())
    scheduler = optim.lr_scheduler.ExponentialLR(optimizer=optimizer, gamma=0.95)
    logout('Done')

    best_iou = 0.0
    for epoch in range(args.epochs):
        st = time.time()
        logout('Beginning Epoch {:02d}'.format(epoch))
        train_sampler.set_epoch(epoch=epoch)
        train_results = train(epoch, net, device, train_data, optimizer, args.batches_per_epoch, vis=args.vis)
        scheduler.step()

        # Run Validation
        if args.local_rank == 0:
            logging.info('Validating...')
            test_results = validate(net, device, val_data, args.val_batches)

        if args.local_rank == 0:
            # Log training losses to tensorboard
            tb.add_scalar('loss/train_loss', train_results['loss'], epoch)
            for n, l in train_results['losses'].items():
                tb.add_scalar('train_loss/' + n, l, epoch)
            logging.info('%d/%d = %f' % (test_results['correct'], test_results['correct'] + test_results['failed'],
                                         test_results['correct'] / (test_results['correct'] + test_results['failed'])))

            logging.info('%d/%d = %f' % (test_results['correct'], test_results['correct'] + test_results['failed'],
                                         test_results['correct'] / (test_results['correct'] + test_results['failed'])))

            # Log validation results to tensorbaord
            tb.add_scalar('loss/IOU', test_results['correct'] / (test_results['correct'] + test_results['failed']), epoch)
            tb.add_scalar('loss/val_loss', test_results['loss'], epoch)
            for n, l in test_results['losses'].items():
                tb.add_scalar('val_loss/' + n, l, epoch)

        # Save best performing network
        iou = test_results['correct'] / (test_results['correct'] + test_results['failed'])
        if iou > best_iou or epoch == 0 or (epoch % 10) == 0:
            if args.local_rank == 0:
                torch.save(net, os.path.join(save_folder, 'epoch_%02d_iou_%0.2f' % (epoch, iou)))
                torch.save(net.state_dict(), os.path.join(save_folder, 'epoch_%02d_iou_%0.2f_statedict.pt' % (epoch, iou)))
            best_iou = iou
        et = time.time()
        logout(f'Epoch-{epoch} took time {et - st} s')


if __name__ == '__main__':
    """
    Usage:
    python -m torch.distributed.launch --nproc_per_node 2 \
        train_ggcnn_multi_gpus.py --network ggcnn \
        --dataset graspnet1b \
        --dataset-path /share/home/ryan/datasets/graspnet-1billion/ \
        --epochs 2 \
        --batch-size 16 \
        --num-workers 20 \
        --description ggcnn3_gn1b_o384_multi \
        --output-size 384 
        --scale 1 \
        --devices 0,1
    """
    args = parse_args()
    run()