train.py

#!/usr/bin/env python3
"""
Main script for running training experiments.

Author: Luigi Celona
Copyright (c) 2026 Luigi Celona
License: MIT
"""

import os
import sys
import argparse
import time

import numpy as np
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from torch.utils.tensorboard import SummaryWriter
from torchvision.utils import make_grid

from loss import LossFunction
from dataset import TripletDataSet
from utils import get_training_dataloader, get_test_dataloader, WarmUpLR, \
    get_torch_network, get_dataloader, get_mean_std, AverageMeter, \
    ProgressMeter, Summary, accuracy

seed = 33
torch.manual_seed(seed)
np.random.seed(seed)


classes = ["Drive_Safe", "Sleep", "Yawning", "Talk_Left", "Talk_Right",
           "Text_Left", "Text_Right", "Make_Up", "Look_Left", "Look_Right",
           "Look_Up", "Look_Down", "Smoke_Left", "Smoke_Right", "Smoke_Mouth",
           "Eat_Left", "Eat_Right", "Operate_Radio", "Operate_GPS",
           "Reach_Behind", "Leave_Steering_Wheel", "Talk_to_Passenger"]


class Net(nn.Module):
    def __init__(self, args):
        super(Net, self).__init__()
        self.net = get_torch_network(args)

        if isinstance(self.net, torchvision.models.squeezenet.SqueezeNet):
            self.net.classifier = nn.AdaptiveAvgPool2d(1)
        else:
            self.net.classifier = nn.Identity()
        dummy_data = torch.randn(1,3,224,224)
        n_feat = self.net(dummy_data).shape[1]
        self.queries = nn.Parameter(torch.randn(3, n_feat))
        self.action_classifier = nn.Linear(n_feat, len(classes))
        self.view_classifier = nn.Linear(n_feat, 3)

    def forward(self, x):
        x = self.net(x)
        y_view = self.view_classifier(x)
        weighted_queries = torch.softmax(y_view, 1).unsqueeze(-1) * self.queries[None]
        x_action = torch.mul(x.unsqueeze(1), weighted_queries).sum(1)
        y_action = self.action_classifier(x_action)
        return x, y_view, x_action, y_action


def train(epoch, writer, args):
    batch_time = AverageMeter('Time', ':6.3f')
    data_time = AverageMeter('Data', ':6.3f')
    lossesACE = AverageMeter('LossACE', ':.4f')
    lossesVCE = AverageMeter('LossVCE', ':.4f')
    lossesVC = AverageMeter('LossVC', ':.4f')
    lossesAC = AverageMeter('LossAC', ':.4f')
    top1A = AverageMeter('Acc@1A', ':6.2f')
    top5A = AverageMeter('Acc@5A', ':6.2f')
    top1V = AverageMeter('Acc@1V', ':6.2f')
    lr = AverageMeter('LR', ':4f', Summary.NONE)

    progress = ProgressMeter(
        len(training_loader),
        [batch_time, data_time, lr, lossesACE, lossesVCE,
         lossesVC, lossesAC, top1A, top1V],
        prefix="Epoch: [{}]".format(epoch))

    net.train()

    end = time.time()
    for batch_index, (img_a, img_sa, img_sv, \
                      label_a, label_sa, label_sv, \
                      view_a, view_sa, view_sv) in enumerate(training_loader):
        data_time.update(time.time() - end)

        img_a = img_a.to(device)
        img_sa = img_sa.to(device)
        img_sv = img_sv.to(device)
        label_a = label_a.to(device)
        label_sa = label_sa.to(device)
        label_sv = label_sa.to(device)
        view_a = view_a.to(device)
        view_sa = view_sa.to(device)
        view_sv = view_sv.to(device)

        x_view_a, y_view_a, x_action_a, y_action_a = net(img_a)
        x_view_sa, y_view_sa, x_action_sa, y_action_sa = net(img_sa)
        x_view_sv, y_view_sv, x_action_sv, y_action_sv = net(img_sv)

        optimizer.zero_grad()
        loss_vce = F.cross_entropy(y_view_a, view_a)
        loss_ace = F.cross_entropy(y_action_a, label_a)
        loss_vc = triplet_loss(x_view_a, x_view_sa, x_view_sv)
        loss_ac = triplet_loss(x_action_a, x_action_sa, x_action_sv)
        loss = loss_ace + loss_vce + args.alpha_triplet_action * loss_ac + args.alpha_triplet_view * loss_vc

        lossesACE.update(loss_ace.item(), img_a.size(0))
        lossesAC.update(loss_ac.item(), img_a.size(0))
        lossesVCE.update(loss_vce.item(), img_a.size(0))
        lossesVC.update(loss_vc.item(), img_a.size(0))

        acc = accuracy(y_action_a.detach(), label_a.detach(), topk=(1,5))
        top1A.update(acc[0].item(), img_a.size(0))
        top5A.update(acc[1].item(), img_a.size(0))

        acc = accuracy(y_view_a.detach(), view_a.detach(), topk=(1,))
        top1V.update(acc[0].item(), img_a.size(0))
        
        lr.update(optimizer.param_groups[0]['lr'], img_a.size(0))
        
        loss.backward()
        optimizer.step()

        #######
        if epoch <= args.warm:
            warmup_scheduler.step()
        #######

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if batch_index % 50 == 0:
            progress.display(batch_index + 1)

    writer.add_scalar('lossesACE/train', lossesACE.avg, epoch)
    writer.add_scalar('lossesVCE/train', lossesVCE.avg, epoch)
    writer.add_scalar('lossesAC/train', lossesAC.avg, epoch)
    writer.add_scalar('lossesVC/train', lossesVC.avg, epoch)
    writer.add_scalar('Top1A/train', top1A.avg, epoch)
    writer.add_scalar('Top5A/train', top5A.avg, epoch)
    writer.add_scalar('Top1V/train', top1V.avg, epoch)
    writer.add_image('images',
                     make_grid(torch.cat([img_a[:5], img_sa[:5], img_sv[:5]]),
                               nrow=5,
                               normalize=True
                               ),
                     epoch
                     )

@torch.no_grad()
def validate(epoch, writer):
    lossesACE = AverageMeter('LossACE', ':.4f')
    lossesVCE = AverageMeter('LossVCE', ':.4f')
    lossesVC = AverageMeter('LossVC', ':.4f')
    lossesAC = AverageMeter('LossAC', ':.4f')
    top1A = AverageMeter('Acc@1A', ':6.2f', Summary.AVERAGE)
    top5A = AverageMeter('Acc@5A', ':6.2f', Summary.AVERAGE)
    top1V = AverageMeter('Acc@1V', ':6.2f', Summary.AVERAGE)

    progress = ProgressMeter(
        len(val_loader),
        [lossesACE, lossesVCE, lossesVC, lossesAC, top1A, top1V],
        prefix='Test: ')

    net.eval()
    # validation
    gts = torch.zeros(len(val_loader.dataset))
    for i, (img_a, img_sa, img_sv, \
            label_a, label_sa, label_sv, \
            view_a, view_sa, view_sv) in enumerate(val_loader):

        img_a = img_a.to(device)
        img_sa = img_sa.to(device)
        img_sv = img_sv.to(device)
        label_a = label_a.to(device)
        label_sa = label_sa.to(device)
        label_sv = label_sa.to(device)
        view_a = view_a.to(device)
        view_sa = view_sa.to(device)
        view_sv = view_sv.to(device)

        x_view_a, y_view_a, x_action_a, y_action_a = net(img_a)
        x_view_sa, y_view_sa, x_action_sa, y_action_sa = net(img_sa)
        x_view_sv, y_view_sv, x_action_sv, y_action_sv = net(img_sv)

        loss_vce = F.cross_entropy(y_view_a, view_a)
        loss_ace = F.cross_entropy(y_action_a, label_a)
        loss_vc = triplet_loss(x_view_a, x_view_sa, x_view_sv)
        loss_ac = triplet_loss(x_action_a, x_action_sa, x_action_sv)

        lossesACE.update(loss_ace.item(), img_a.size(0))
        lossesAC.update(loss_ac.item(), img_a.size(0))
        lossesVCE.update(loss_vce.item(), img_a.size(0))
        lossesVC.update(loss_vc.item(), img_a.size(0))

        acc = accuracy(y_action_a.detach(), label_a.detach(), topk=(1,5))
        top1A.update(acc[0].item(), img_a.size(0))
        top5A.update(acc[1].item(), img_a.size(0))

        acc = accuracy(y_view_a.detach(), view_a.detach(), topk=(1,))
        top1V.update(acc[0].item(), img_a.size(0))

    writer.add_scalar('lossesACE/val', lossesACE.avg, epoch)
    writer.add_scalar('lossesVCE/val', lossesVCE.avg, epoch)
    writer.add_scalar('lossesAC/val', lossesAC.avg, epoch)
    writer.add_scalar('lossesVC/val', lossesVC.avg, epoch)
    writer.add_scalar('Top1A/val', top1A.avg, epoch)
    writer.add_scalar('Top5A/val', top5A.avg, epoch)
    writer.add_scalar('Top1V/val', top1V.avg, epoch)

    if args.gpu:
        print('GPU INFO.....')
        print(torch.cuda.memory_summary())
    print('Evaluating Network.....')
    progress.display_summary()
    return top1A.avg


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('-net', type=str, required=True, help='net type')
    parser.add_argument('-gpu', action='store_true', default=True, help='use gpu or not')
    parser.add_argument('-b', type=int, default=64, help='batch size for dataloader')
    parser.add_argument('-warm', type=int, default=2, help='warm up training phase')
    parser.add_argument('-lr', type=float, default=0.1, help='initial learning rate')
    parser.add_argument('-resume', action='store_true', default=False, help='resume training')
    parser.add_argument('-pretrain', default=True, help='whether the pretrain model is used')
    parser.add_argument("-data_path" , required=True, help='path to the dataset')
    parser.add_argument("-dataset", default='pic-day-cam1', type=str)
    parser.add_argument("-num_class", default=22, type=int)
    parser.add_argument("-exp_dir", default="contrastive", type=str)
    parser.add_argument("-alpha_triplet_view", default=1, type=float)
    parser.add_argument("-alpha_triplet_action", default=1, type=float)

    args = parser.parse_args()

    args.exp_dir = args.exp_dir + "_alpha_tv%.2f" % args.alpha_triplet_view + "_alpha_ta%.2f" % args.alpha_triplet_action

    device = torch.device("cuda" if args.gpu else "cpu")

# build the backbone model
    net = Net(args).to(device)

    for name,parameters in net.named_parameters():
        print(name, ':', parameters.size())

    triplet_loss = LossFunction()

# the config for wandb output
    writer = SummaryWriter(args.exp_dir + '/' + args.dataset + '/' + args.net)

# prepare the data
    trainloader, trainloadertxt, valloader, valloadertxt, *_ = get_dataloader(args.data_path, args.dataset)
    mean, std = get_mean_std(args.dataset)
    train_datasets = TripletDataSet(trainloader, trainloadertxt, mean, std, flag ='train')#get data
    val_datasets = TripletDataSet(valloader, valloadertxt, mean, std, flag = 'val')#
    training_loader = get_training_dataloader(
        dataset = train_datasets,
        num_workers=0,
        batch_size=args.b,
        shuffle=True  
    )
    val_loader = get_test_dataloader(
        dataset = val_datasets,
        num_workers=0,
        batch_size=1,
        shuffle=False  
    )

    iter_per_epoch = len(training_loader)
    print("Total params ", sum(p.numel() for p in net.parameters()))
    print("Trainable params ", sum(p.numel() for p in net.parameters() if p.requires_grad))

# set detail of the checkpoint storage path
    checkpoint_path = os.path.join(args.exp_dir, args.dataset, args.net)
    if not os.path.exists(checkpoint_path):
        os.makedirs(checkpoint_path)
    checkpoint_path = os.path.join(checkpoint_path, '{net}-{epoch}-{type}.pth')

    start_epoch = 1

# prepare the optimizer
    optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)
    train_scheduler = optim.lr_scheduler.MultiStepLR(optimizer,
                                                     milestones=[10, 30],
                                                     gamma=0.5) #learning rate decay
    warmup_scheduler = WarmUpLR(optimizer, iter_per_epoch * args.warm)

    if args.resume:
        print("Resume last checkpoint")
        ckpt_list = np.asarray([p for p in os.listdir(os.path.dirname(checkpoint_path)) if '.pth' in p])
        resume_epoch = max([int(c.split('-')[1]) for c in ckpt_list])
        idx = np.asarray([str(resume_epoch) in p for p in ckpt_list])
        ckpt_fn = os.path.join(os.path.dirname(checkpoint_path), ckpt_list[idx].item())
        start_epoch = resume_epoch + 1
        ckpt = torch.load(ckpt_fn)
        net.load_state_dict(ckpt['model'])
        optimizer.load_state_dict(ckpt['optimizer'])
        train_scheduler.load_state_dict(ckpt['scheduler'])

# start training the model
    best_acc = 0.0
    for epoch in range(start_epoch, 21): #51):
        if epoch > args.warm:
            train_scheduler.step()
        if args.resume:
            if epoch <= resume_epoch:
                continue
        train(epoch, writer, args)
        acc = validate(epoch, writer)

        # save best performance model after learning rate decay to 0.01
        ckpt = {'model': net.state_dict(),
                'optimizer': optimizer.state_dict(),
                'scheduler': train_scheduler.state_dict()
                }

        if best_acc < acc:
            print("best model! save...")
            torch.save(ckpt,
                       checkpoint_path.format(net=args.net,
                                              epoch=epoch,
                                              type='best'))
            best_acc = acc
            continue

        if not epoch % 5:
            torch.save(ckpt,
                       checkpoint_path.format(net=args.net,
                                              epoch=epoch,
                                              type='regular'))