train.py

import argparse
import os
import sys
import math
import torch
import random
import logging
import numpy as np
from tqdm.auto import tqdm
import torch.optim as optim
from datetime import datetime
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

# Add sggnet to path
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, ROOT_DIR)

from sggnet.models.hggqnet import HGGQNet
from sggnet.models.loss import get_loss
from sggnet.dataset.graspnet_dataset import GraspNetDataset, collate_fn

parser = argparse.ArgumentParser()
parser.add_argument('--dataset_root', required=True, help='Dataset root')
parser.add_argument('--camera', required=True, help='Camera split [realsense/kinect]')
parser.add_argument('--checkpoint_path', default=None, help='Model checkpoint path [default: None]')
parser.add_argument('--num_points', type=int, default=10000, help='Points Number [default: 10000]')
parser.add_argument('--max_epoch', type=int, default=11, help='Epoch to run [default: 11]')
parser.add_argument('--batch_size', type=int, default=1, help='Batch Size during training [default: 1]')
parser.add_argument('--mode', default='optimize_parameter', help='Mode')
cfgs = parser.parse_args()

random.seed(42)
np.random.seed(42)
torch.manual_seed(42)

LOG_DIR = os.path.join('log')
os.makedirs(LOG_DIR, exist_ok=True)

logging.basicConfig(
    filename=os.path.join(LOG_DIR, 'train.log'),
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger()


def log_string(out_str, level=logging.INFO):
    if level == logging.DEBUG:
        logger.debug(out_str)
    elif level == logging.WARNING:
        logger.warning(out_str)
    elif level == logging.ERROR:
        logger.error(out_str)
    else:
        logger.info(out_str)


def evaluate_one_epoch(epoch, model, dataloader, writer, device):
    model.eval()
    total_loss = 0.0
    total_node_score_loss = 0.0
    total_grasp_score_loss = 0.0
    num_batches = len(dataloader)

    log_string(f"Starting Evaluation Epoch {epoch}...")

    with tqdm(dataloader, total=num_batches, desc=f'Eval Epoch {epoch}') as t:
        for batch_idx, batch_data in enumerate(t):
            with torch.no_grad():
                grasp_configs_tensor, grasp_config_scores_tensor, graph, graph_indices = batch_data
                grasp_configs_tensor = grasp_configs_tensor.to(device)
                grasp_config_scores_tensor = grasp_config_scores_tensor.to(device)
                graph = graph.to(device)
                graph_indices = graph_indices.to(device)

                predicted_scores, predicted_center_scores, predicted_fingers_scores = model(
                    grasp_configs_tensor, graph, graph_indices)

                # Compute loss
                loss, center_score_loss, grasp_score_loss, fingers_score_loss = get_loss(
                    predicted_scores, predicted_center_scores, predicted_fingers_scores,
                    grasp_config_scores_tensor, graph)
                
                total_loss += loss.item()
                total_node_score_loss += center_score_loss.item()
                total_grasp_score_loss += grasp_score_loss.item()

                t.set_postfix(
                    Loss=f"{total_loss / (batch_idx + 1):.4f}",
                    NodeLoss=f"{total_node_score_loss / (batch_idx + 1):.4f}",
                    GraspLoss=f"{total_grasp_score_loss / (batch_idx + 1):.4f}"
                )

                iteration = epoch * num_batches + batch_idx
                writer.add_scalar('ValLoss', loss.item(), iteration)
                writer.add_scalar('ValLoss/NodeScore', center_score_loss.item(), iteration)
                writer.add_scalar('ValLoss/GraspScore', grasp_score_loss.item(), iteration)

    avg_loss = total_loss / num_batches
    avg_node_score_loss = total_node_score_loss / num_batches
    avg_grasp_score_loss = total_grasp_score_loss / num_batches

    writer.add_scalar('ValLoss/Avg', avg_loss, epoch)
    writer.add_scalar('ValLoss/NodeScore/Avg', avg_node_score_loss, epoch)
    writer.add_scalar('ValLoss/GraspScore/Avg', avg_grasp_score_loss, epoch)

    log_string(f"[Epoch {epoch}] "
               f"Loss: {avg_loss:.4f}, Node Loss: {avg_node_score_loss:.4f}, Grasp Loss: {avg_grasp_score_loss:.4f}")

    return avg_loss


def train(start_epoch, optimizer, scheduler, model, train_dataloader, test_dataloader, writer,
          checkpoint_dir, device, mode=None):
    log_string("Starting training process...")

    for epoch in range(start_epoch, cfgs.max_epoch):
        epoch_start_time = datetime.now()

        mean_loss = train_one_epoch(epoch, model, optimizer, train_dataloader, writer, checkpoint_dir, device)

        mean_val_loss = 0
        if mode != "optimize_parameter" and test_dataloader is not None:
            mean_val_loss = evaluate_one_epoch(epoch, model, test_dataloader, writer, device)

        if scheduler is not None:
            if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                scheduler.step(mean_loss if mode == "optimize_parameter" else mean_val_loss)
            else:
                scheduler.step()

        # Save epoch checkpoint
        save_dict = {
            'loss': mean_loss,
            'val_loss': mean_val_loss,
            'epoch': epoch + 1,
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict() if scheduler else None,
            'model_state_dict': model.state_dict()
        }
        checkpoint_path = os.path.join(checkpoint_dir, f'checkpoint_epoch{epoch}.tar')
        torch.save(save_dict, checkpoint_path)
        log_string(f"Saved epoch checkpoint for epoch {epoch}.")

        epoch_end_time = datetime.now()
        log_string(f"Epoch {epoch} completed in {epoch_end_time - epoch_start_time}.")

    log_string("Training completed.")


def train_one_epoch(epoch, model, optimizer, train_dataloader, writer, checkpoint_dir, device):
    model.train()
    running_loss = 0.0
    running_center_score_loss = 0.0
    running_fingers_score_loss = 0.0
    running_grasp_score_loss = 0.0
    num_batches = len(train_dataloader)

    log_string(f"Starting Training Epoch {epoch}...")

    with tqdm(train_dataloader, total=num_batches, desc=f'Train Epoch {epoch}') as t:
        for batch_idx, batch_data in enumerate(t):
            grasp_configs_tensor, grasp_config_scores_tensor, graph, graph_indices = batch_data
            grasp_configs_tensor = grasp_configs_tensor.to(device)
            grasp_config_scores_tensor = grasp_config_scores_tensor.to(device)
            graph = graph.to(device)
            graph_indices = graph_indices.to(device)

            optimizer.zero_grad()

            predicted_scores, predicted_center_scores, predicted_fingers_scores = model(
                grasp_configs_tensor, graph, graph_indices)

            loss, center_score_loss, grasp_score_loss, fingers_score_loss = get_loss(
                predicted_scores, predicted_center_scores, predicted_fingers_scores,
                grasp_config_scores_tensor, graph)

            if not math.isnan(loss.item()):
                loss.backward()
                optimizer.step()

                running_loss += loss.item()
                running_center_score_loss += center_score_loss.item()
                running_fingers_score_loss += fingers_score_loss.item()
                running_grasp_score_loss += grasp_score_loss.item()

                iteration = epoch * num_batches + batch_idx
                writer.add_scalar('Loss', loss.item(), iteration)
                writer.add_scalar('Loss/CenterScore', center_score_loss.item(), iteration)
                writer.add_scalar('Loss/FingersScore', fingers_score_loss.item(), iteration)
                writer.add_scalar('Loss/GraspScore', grasp_score_loss.item(), iteration)

                t.set_postfix(
                    Loss=f"{running_loss / (batch_idx + 1):.4f}",
                    CenterLoss=f"{running_center_score_loss / (batch_idx + 1):.4f}",
                    FingersLoss=f"{running_fingers_score_loss / (batch_idx + 1):.4f}",
                    GraspLoss=f"{running_grasp_score_loss / (batch_idx + 1):.4f}"
                )

            torch.cuda.empty_cache()

    avg_loss = running_loss / num_batches
    avg_center_score_loss = running_center_score_loss / num_batches
    avg_fingers_score_loss = running_fingers_score_loss / num_batches
    avg_grasp_score_loss = running_grasp_score_loss / num_batches

    writer.add_scalar('Loss/Avg', avg_loss, epoch)
    writer.add_scalar('Loss/CenterScore/Avg', avg_center_score_loss, epoch)
    writer.add_scalar('Loss/FingersScore/Avg', avg_fingers_score_loss, epoch)
    writer.add_scalar('Loss/GraspScore/Avg', avg_grasp_score_loss, epoch)

    log_string(f"[Epoch {epoch}] "
               f"Loss: {avg_loss:.4f}, Center Loss: {avg_center_score_loss:.4f}, "
               f"Fingers Loss: {avg_fingers_score_loss:.4f}, Grasp Loss: {avg_grasp_score_loss:.4f}")

    return avg_loss


def worker_init_fn(worker_id):
    torch.utils.data.get_worker_info().dataset.load_grasp_labels()


if __name__ == '__main__':
    # Configure logging dynamically
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    log_string(f"Using device: {device}")

    # Load datasets
    train_dataset = GraspNetDataset(cfgs.dataset_root, camera=cfgs.camera, split='train',
                                    num_points=cfgs.num_points, remove_outlier=True)
    train_dataset.load_grasp_labels()
    train_dataloader = DataLoader(train_dataset, batch_size=cfgs.batch_size, shuffle=True,
                                  collate_fn=collate_fn, worker_init_fn=worker_init_fn, num_workers=8,
                                  persistent_workers=True)

    test_dataloader = None
    if cfgs.mode != "optimize_parameter":
        test_dataset = GraspNetDataset(cfgs.dataset_root, camera=cfgs.camera, split='test_seen',
                                       num_points=cfgs.num_points, remove_outlier=True)
        test_dataset.load_grasp_labels()
        test_dataloader = DataLoader(test_dataset, batch_size=cfgs.batch_size, shuffle=False,
                                     collate_fn=collate_fn)

    # Create model
    model = HGGQNet()
    model.to(device)

    # Load checkpoint if provided
    start_epoch = 0
    if cfgs.checkpoint_path and os.path.exists(cfgs.checkpoint_path):
        checkpoint = torch.load(cfgs.checkpoint_path)
        model.load_state_dict(checkpoint['model_state_dict'])
        start_epoch = checkpoint.get('epoch', 0)
        log_string(f"Loaded checkpoint from {cfgs.checkpoint_path} (epoch: {start_epoch})")

    # Setup optimizer and scheduler
    optimizer = optim.AdamW(model.parameters(), lr=0.01, weight_decay=0.01)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)

    # Setup logging and checkpoint directories
    root_log_dir = 'log/tensorboard'
    os.makedirs(root_log_dir, exist_ok=True)
    config_name = "adamw_lr0.01_wd0.01_cosine_annealing_lr_Tmax10"
    config_log_dir = os.path.join(root_log_dir, config_name)
    os.makedirs(config_log_dir, exist_ok=True)

    checkpoint_dir = os.path.join(config_log_dir, 'checkpoints')
    os.makedirs(checkpoint_dir, exist_ok=True)

    # Set up TensorBoard writer
    writer = SummaryWriter(config_log_dir)

    log_string(f"Starting training with configuration: {config_name}")

    # Train
    train(start_epoch, optimizer, scheduler, model, train_dataloader, test_dataloader, writer, 
          checkpoint_dir, device, mode=cfgs.mode)

    log_string("Training completed.")