training/scripts/train.py

#!/usr/bin/env python
#
# file: $ISIP_EXP/SOGMP/scripts/train.py
#
# revision history: xzt
#  20220824 (TE): first version
#
# usage:
#  python train.py mdir train_data val_data
#
# arguments:
#  mdir: the directory where the output model is stored
#  train_data: the directory of training data
#  val_data: the directory of valiation data
#
# This script trains a Semantic CNN model
#------------------------------------------------------------------------------

# import pytorch modules
#
import torch
import torch.nn as nn
from torch.optim import Adam
from tqdm import tqdm

# visualize:
from tensorboardX import SummaryWriter
import numpy as np

# import the model and all of its variables/functions
#
from model import *

# import modules
#
import sys
import os


#-----------------------------------------------------------------------------
#
# global variables are listed here
#
#-----------------------------------------------------------------------------

# general global values
#
model_dir = './model/semantic_cnn_model.pth'  # the path of model storage 
NUM_ARGS = 3
NUM_EPOCHS = 4000 
BATCH_SIZE = 64 
LEARNING_RATE = "lr"
BETAS = "betas"
EPS = "eps"
WEIGHT_DECAY = "weight_decay"

# for reproducibility, we seed the rng
#
set_seed(SEED1)       

# adjust_learning_rate
#　
def adjust_learning_rate(optimizer, epoch):
    lr = 1e-3
    if epoch > 40:
        lr = 2e-4
    if epoch > 2000:
        lr = 2e-5
    if epoch > 21000:
        lr = 1e-5
    if epoch > 32984:
        lr = 1e-6
    if epoch > 48000:
       # lr = 5e-8
       lr = lr * (0.1 ** (epoch // 110000))
    #  if epoch > 8300:
    #      lr = 1e-9
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr



# train function:
def train(model, dataloader, dataset, device, optimizer, criterion, epoch, epochs):
    ################################## Train #####################################
    # Set model to training mode
    model.train()  
    # for each batch in increments of batch size
    #
    running_loss = 0
    counter = 0
    # get the number of batches (ceiling of train_data/batch_size):
    num_batches = int(len(dataset)/dataloader.batch_size)
    for i, batch in tqdm(enumerate(dataloader), total=num_batches):
    #for i, batch in enumerate(dataloader, 0):
        counter += 1
        # collect the samples as a batch:
        scan_maps = batch['scan_map']
        scan_maps = scan_maps.to(device)
        semantic_maps = batch['semantic_map']
        semantic_maps = semantic_maps.to(device)
        sub_goals = batch['sub_goal']
        sub_goals = sub_goals.to(device)
        velocities = batch['velocity']
        velocities = velocities.to(device)

        # set all gradients to 0:
        optimizer.zero_grad()
        # feed the network the batch
        #
        
        output = model(scan_maps, semantic_maps, sub_goals)
        #writer.add_graph(model,[batch_ped_pos_t, batch_scan_t, batch_goal_t])    
        # get the loss
        #
        # loss = criterion(output, velocities)

        # ---------------------------
        # Mask zero-velocity samples
        # ---------------------------
        mask = (velocities != 0).any(dim=1)  # (B,)

        if mask.sum() == 0:
            loss = output.sum() * 0   # safe zero loss
        else:
            loss = criterion(output[mask], velocities[mask])

        # perform back propagation:
        loss.backward(torch.ones_like(loss))
        optimizer.step()
        # get the loss:
        # multiple GPUs:
        if torch.cuda.device_count() > 1:
            loss = loss.mean()  

        running_loss += loss.item()
        
        # display informational message
        #
        if(i % 1280 == 0):
            print('Epoch [{}/{}], Step[{}/{}], Loss: {:.4f}'
                    .format(epoch, epochs, i + 1, num_batches, loss.item()))

    train_loss = running_loss / len(dataset) #counter 

    return train_loss

# validate function:
def validate(model, dataloader, dataset, device, criterion):
    ################################## Train #####################################
    # set model to evaluation mode:
    model.eval()
    # for each batch in increments of batch size
    #
    running_loss = 0
    counter = 0
    # get the number of batches (ceiling of train_data/batch_size):
    num_batches = int(len(dataset)/dataloader.batch_size)
    for i, batch in tqdm(enumerate(dataloader), total=num_batches):
    #for i, batch in enumerate(dataloader, 0):
        counter += 1
        # collect the samples as a batch:
        scan_maps = batch['scan_map']
        scan_maps = scan_maps.to(device)

        semantic_maps = batch['semantic_map']
        semantic_maps = semantic_maps.to(device)
 
        sub_goals = batch['sub_goal']
        sub_goals = sub_goals.to(device)
        velocities = batch['velocity']
        velocities = velocities.to(device)

        # feed the network the batch
        #
        output = model(scan_maps, semantic_maps, sub_goals)
        #writer.add_graph(model,[batch_ped_pos_t, batch_scan_t, batch_goal_t])    
        # get the loss
        #
        # loss = criterion(output, velocities)
        # ---------------------------
        # Mask zero-velocity samples
        # ---------------------------
        mask = (velocities != 0).any(dim=1)  # (B,)

        if mask.sum() == 0:
            loss = output.sum() * 0   # safe zero loss
        else:
            loss = criterion(output[mask], velocities[mask])
            
        # get the loss:
        # multiple GPUs:
        if torch.cuda.device_count() > 1:
            loss = loss.mean()  

        running_loss += loss.item()

    val_loss = running_loss / len(dataset) #counter 

    return val_loss

#------------------------------------------------------------------------------
#
# the main program starts here
#
#------------------------------------------------------------------------------

# function: main
#
# arguments: none
#
# return: none
#
# This method is the main function.
#
def main(argv):

    # ensure we have the correct amount of arguments
    #
    #global cur_batch_win
    if(len(argv) != NUM_ARGS):
        print("usage: python nedc_train_mdl.py [MDL_PATH] [TRAIN_PATH] [DEV_PATH]")
        exit(-1)

    # define local variables
    #
    mdl_path = argv[0]
    pTrain = argv[1]
    pDev = argv[2]

    # get the output directory name
    #
    odir = os.path.dirname(mdl_path)

    # if the odir doesn't exits, we make it
    #
    if not os.path.exists(odir):
        os.makedirs(odir)

    # set the device to use GPU if available
    #
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    ### train:
    print('...Start reading data...')
    # get array of the data
    # data: [[0, 1, ... 26], [27, 28, ...] ...]
    # labels: [0, 0, 1, ...]
    #
    #[ped_pos_t, scan_t, goal_t, vel_t] = get_data(pTrain)
    train_dataset = NavDataset(pTrain, 'train')
    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, \
                                                   shuffle=True, drop_last=True, pin_memory=True)
    #train_data = train_data - np.mean(train_data, axis=0)
    
    ### dev:

    # get array of the data
    # data: [[0, 1, ... 26], [27, 28, ...] ...]
    # labels: [0, 0, 1, ...]
    #
    #[ped_pos_d, scan_d, goal_d, vel_d] = get_data(pDev)
    dev_dataset = NavDataset(pDev, 'dev')
    dev_dataloader = torch.utils.data.DataLoader(dev_dataset, batch_size=BATCH_SIZE, \
                                                   shuffle=True, drop_last=True, pin_memory=True)
    #dev_data = dev_data - np.mean(dev_data, axis=0)
    print('...Finish reading data...')

    # instantiate a model
    #
    model = SemanticCNN(Bottleneck, [2, 1, 1])

    # moves the model to device (cpu in our case so no change)
    #
    model.to(device)

    # set the adam optimizer parameters
    #
    opt_params = { LEARNING_RATE: 0.001,
                   BETAS: (.9,0.999),
                   EPS: 1e-08,
                   WEIGHT_DECAY: .001 }

    # set the loss and optimizer
    #
    criterion = nn.MSELoss(reduction='sum')
    criterion.to(device)

    # create an optimizer, and pass the model params to it
    #
    optimizer = Adam(model.parameters(), **opt_params)

    # get the number of epochs to train on
    #
    epochs = NUM_EPOCHS
    
    # if there are trained models, continue training:
    if os.path.exists(mdl_path):
        checkpoint = torch.load(mdl_path)
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        start_epoch = checkpoint['epoch']
        print('Load epoch {} success'.format(start_epoch))
    else:
        start_epoch = 0
        print('No trained models, restart training')

    # multiple GPUs:
    if torch.cuda.device_count() > 1:
        print("Let's use 2 of total", torch.cuda.device_count(), "GPUs!")
        # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
        model = nn.DataParallel(model) #, device_ids=[0, 1])

    # moves the model to device (cpu in our case so no change)
    #
    model.to(device)

    # tensorboard writer:
    writer = SummaryWriter('runs')

    # for each epoch
    #
    #loss_train = []
    #loss_vector = []
    epoch_num = 0
    for epoch in range(start_epoch+1, epochs):

        # adjust learning rate:
        adjust_learning_rate(optimizer, epoch)
        ################################## Train #####################################
        # for each batch in increments of batch size
        #
        train_epoch_loss = train(
            model, train_dataloader, train_dataset, device, optimizer, criterion, epoch, epochs
        )
        
        ################################## Test #####################################
        valid_epoch_loss = validate(
            model, dev_dataloader, dev_dataset, device, criterion
        )

        # log the epoch loss
        writer.add_scalar('training loss',
                        train_epoch_loss,
                        epoch)
        writer.add_scalar('validation loss',
                        valid_epoch_loss,
                        epoch)

        print('Train set: Average loss: {:.4f}'.format(train_epoch_loss))
        print('Validation set: Average loss: {:.4f}'.format(valid_epoch_loss))

        # save the model
        #
        if(epoch % 50 == 0):
            if torch.cuda.device_count() > 1: # multiple GPUS: 
                state = {'model':model.module.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}
            else:
                state = {'model':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}
            path='./model/model' + str(epoch) +'.pth'
            torch.save(state, path)
        
        epoch_num = epoch

    # save the final model
    if torch.cuda.device_count() > 1: # multiple GPUS: 
        state = {'model':model.module.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch_num}
    else:
        state = {'model':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch_num}
    torch.save(state, mdl_path)

    # exit gracefully
    #

    return True
#
# end of function


# begin gracefully
#
if __name__ == '__main__':
    main(sys.argv[1:])
#
# end of file