Upload train.py with huggingface_hub

train.py

@@ -0,0 +1,613 @@
+import numpy as np
+import os
+import random
+import shutil
+import time
+import warnings
+from enum import Enum
+import pandas as pd
+import torch
+import logging
+
+import torch.nn.parallel
+import torch.optim
+import torch.utils.data
+import torch.utils.data.distributed
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import torchmetrics
+from torch.utils.tensorboard import SummaryWriter
+
+from configs.args_base import get_args
+
+from utils.losses import build_loss_function
+from utils.optimizer import build_optimizer
+from utils.lr_scheduler import build_scheduler
+from utils.logger import build_logger
+from utils.misc import setup_seed, reduce_tensor, save_checkpoint
+from data import build_dataloader
+from models.MIQA_base import get_torch_model, get_timm_model
+from models.RA_MIQA import RegionVisionTransformer
+
+best_srcc = best_plcc = best_klcc = 0.
+
+def main(args):
+    if args.seed is not None:
+
+        setup_seed(args.seed)
+
+        warnings.warn('You have chosen to seed training. '
+                      'This will turn on the CUDNN deterministic setting, '
+                      'which can slow down your training considerably! '
+                      'You may see unexpected behavior when restarting '
+                      'from checkpoints.')
+
+    if args.gpu is not None:
+        warnings.warn('You have chosen a specific GPU. This will completely '
+                      'disable data parallelism.')
+
+    if args.dist_url == "env://" and args.world_size == -1:
+        args.world_size = int(os.environ["WORLD_SIZE"])
+
+    args.distributed = args.world_size > 1 or args.multiprocessing_distributed
+
+    if torch.cuda.is_available():
+        ngpus_per_node = torch.cuda.device_count()
+        if ngpus_per_node == 1 and args.dist_backend == "nccl":
+            warnings.warn(
+                "nccl backend >=2.5 requires GPU count>1, see https://github.com/NVIDIA/nccl/issues/103 perhaps use 'gloo'")
+    else:
+        ngpus_per_node = 1
+
+    if args.multiprocessing_distributed:
+        # Since we have ngpus_per_node processes per node, the total world_size
+        # needs to be adjusted accordingly
+        args.world_size = ngpus_per_node * args.world_size
+        # Use torch.multiprocessing.spawn to launch distributed processes: the
+        # main_worker process function
+        mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))
+    else:
+        # Simply call main_worker function
+        main_worker(args.gpu, ngpus_per_node, args)
+
+
+def main_worker(gpu, ngpus_per_node, args):
+    global best_srcc, best_plcc, best_klcc
+    args.gpu = gpu
+    args.ngpus_per_node = ngpus_per_node
+    if args.distributed:
+        if args.dist_url == "env://" and args.rank == -1:
+            args.rank = int(os.environ["RANK"])
+        if args.multiprocessing_distributed:
+            # For multiprocessing distributed training, rank needs to be the
+            # global rank among all the processes
+            args.rank = args.rank * ngpus_per_node + gpu
+        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
+                                world_size=args.world_size, rank=args.rank)
+
+    logger = build_logger(
+        output_dir=args.output_dir,
+        log_file='{}_train.log'.format(args.run_name),
+        rank=args.rank if args.distributed else None,
+        level=logging.INFO,
+        console_level=logging.INFO if args.rank in [0, None] else logging.WARNING,
+        file_level=logging.INFO
+    )
+
+    if args.gpu is not None:
+        logger.info("Use GPU: {} for training".format(args.gpu))
+
+    # create model
+    if args.arch.startswith('RA_'):
+        model = RegionVisionTransformer(
+            base_model_name = 'vit_small_patch16_224',
+            pretrained = True,
+            mmseg_config_path = 'models/model_configs/fcn_sere-small_finetuned_fp16_8x32_224x224_3600_imagenets919.py',
+            checkpoint_path = 'models/checkpoints/sere_finetuned_vit_small_ep100.pth',
+            auto_download = True,
+            force_download = False
+            )
+    else:
+        try:
+            logger.info(f"Loading model form torchvision {args.arch}")
+            model = get_torch_model(model_name=args.arch, pretrained=args.pretrained, num_classes=1)
+        except:
+            logger.info(f"Loading model form timm {args.arch}")
+            model = get_timm_model(model_name=args.arch, pretrained=args.pretrained, num_classes=1)
+
+    if not torch.cuda.is_available() and not torch.backends.mps.is_available():
+        logger.info('using CPU, this will be slow')
+
+    elif args.distributed:
+        # For multiprocessing distributed, DistributedDataParallel constructor
+        # should always set the single device scope, otherwise,
+        # DistributedDataParallel will use all available devices.
+        if torch.cuda.is_available():
+            if args.gpu is not None:
+                torch.cuda.set_device(args.gpu)
+                model.cuda(args.gpu)
+                # When using a single GPU per process and per
+                # DistributedDataParallel, we need to divide the batch size
+                # ourselves based on the total number of GPUs of the current node.
+                args.batch_size = int(args.batch_size / ngpus_per_node)
+                args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node)
+                model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
+            else:
+                model.cuda()
+                # DistributedDataParallel will divide and allocate batch_size to all
+                # available GPUs if device_ids are not set
+                model = torch.nn.parallel.DistributedDataParallel(model)
+
+    elif args.gpu is not None and torch.cuda.is_available():
+        torch.cuda.set_device(args.gpu)
+        model = model.cuda(args.gpu)
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps")
+        model = model.to(device)
+    else:
+        # DataParallel will divide and allocate batch_size to all available GPUs
+        if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
+            model.features = torch.nn.DataParallel(model.features)
+            model.cuda()
+        else:
+            model = torch.nn.DataParallel(model).cuda()
+
+    if torch.cuda.is_available():
+        if args.gpu:
+            device = torch.device('cuda:{}'.format(args.gpu))
+        else:
+            device = torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
+
+    # Data loading
+    train_dataset, val_dataset = build_dataloader.build_dataset(args)
+
+    if args.distributed:
+        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
+        val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False, drop_last=False)
+    else:
+        train_sampler = None
+        val_sampler = None
+
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
+        num_workers=args.workers, pin_memory=True, sampler=train_sampler)
+
+    val_loader = torch.utils.data.DataLoader(
+        val_dataset, batch_size=args.batch_size, shuffle=False,
+        num_workers=args.workers, pin_memory=True, sampler=val_sampler)
+
+    # define loss function (criterion), optimizer, and learning rate scheduler
+    criterion = build_loss_function(loss_name=args.loss_name)
+    optimizer = build_optimizer(args, model)
+    scheduler = build_scheduler(args, optimizer, len(train_loader))
+
+    # optionally resume from a checkpoint
+    if args.resume:
+        if os.path.isfile(args.resume):
+            logger.info("=> loading checkpoint '{}'".format(args.resume))
+            if args.gpu is None:
+                checkpoint = torch.load(args.resume)
+            elif torch.cuda.is_available():
+                # Map model to be loaded to specified single gpu.
+                loc = 'cuda:{}'.format(args.gpu)
+                checkpoint = torch.load(args.resume, map_location=loc)
+            args.start_epoch = checkpoint['epoch']
+            best_srcc = checkpoint['best_srcc']
+            if args.gpu is not None:
+                # best_acc1 may be from a checkpoint from a different GPU
+                best_srcc = best_srcc.to(args.gpu)
+            model.load_state_dict(checkpoint['state_dict'])
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            scheduler.load_state_dict(checkpoint['scheduler'])
+            logger.info("=> loaded checkpoint '{}' (epoch {})"
+                  .format(args.resume, checkpoint['epoch']))
+        else:
+            logger.info("=> no checkpoint found at '{}'".format(args.resume))
+
+    # evaluate on validation set
+    if args.eval_only:
+        validate(val_loader, model, criterion, args)
+        return
+
+    writer = SummaryWriter(log_dir=os.path.join(args.output_dir, 'tensorboard_logs', args.run_name))
+
+    for epoch in range(args.start_epoch, args.epochs+1):
+        if args.distributed:
+            train_sampler.set_epoch(epoch)
+
+        # train for one epoch
+        best_srcc, best_plcc, best_klcc = train_one_epoch(train_loader, model, criterion, optimizer, scheduler, epoch, device, args, val_loader, writer, logger)
+
+    writer.close()
+
+    logger.info('################# Training Finished ##################')
+    logger.info(f"Best SRCC: {best_srcc}, Best PLCC: {best_plcc}, Best KLCC: {best_klcc}")
+    logger.info('######################################################')
+
+
+def train_one_epoch(train_loader, model, criterion, optimizer, scheduler, epoch, device, args, val_loader, writer, logger):
+    global best_srcc, best_plcc, best_klcc
+    model.train()
+
+    batch_time = AverageMeter('Time', ':6.3f')
+    data_time = AverageMeter('Data', ':6.3f')
+    losses = AverageMeter('Loss', ':.4e')
+    srcc = AverageMeter('SRCC', ':6.4f')
+    plcc = AverageMeter('PLCC', ':6.4f')
+    klcc = AverageMeter('KLCC', ':6.4f')
+    # mse = AverageMeter('MSE', ':6.4f')
+
+    progress = ProgressMeter(
+        len(train_loader),
+        [batch_time, data_time, losses, srcc, plcc, klcc],
+        prefix="Epoch: [{}]".format(epoch))
+
+    validate_freq = len(train_loader) // args.validate_num
+
+    end = time.time()
+    global_step = epoch * len(train_loader)
+
+    for i, batch in enumerate(train_loader):
+        data_time.update(time.time() - end)
+
+        # images = batch['image'].cuda(args.gpu, non_blocking=True)
+        # target = batch['label'].cuda(args.gpu, non_blocking=True).view(-1)
+
+        image_cropped = batch['image_cropped'].cuda(args.gpu, non_blocking=True)
+        target = batch['label'].cuda(args.gpu, non_blocking=True).view(-1)
+
+        if 'image_resized' in batch:
+            image_resized = batch['image_resized'].cuda(args.gpu, non_blocking=True)
+            output = model(image_cropped, image_resized).view(-1)
+        else:
+            output = model(image_cropped).view(-1)
+
+        target_len = target.size(0)
+        train_loss = criterion(output, target)
+
+        # Calculate metrics during training sessions
+        srcc_train = torchmetrics.functional.spearman_corrcoef(output, target).item()
+        plcc_train = torchmetrics.functional.pearson_corrcoef(output, target).item()
+        klcc_train = torchmetrics.functional.kendall_rank_corrcoef(output, target).item()
+
+        # Update Metrics
+        losses.update(train_loss.item(), target_len)
+        srcc.update(srcc_train, target_len)
+        plcc.update(plcc_train, target_len)
+        klcc.update(klcc_train, target_len)
+
+        # Add training loss to the writer
+        writer.add_scalars('Loss', {
+            'train': train_loss.item()
+        }, global_step + i)
+
+        optimizer.zero_grad()
+        train_loss.backward()
+        optimizer.step()
+
+        # if scheduler is not None:
+        scheduler.step_update(global_step + i)
+
+        # Record the current learning rate
+        if args.rank == 0:
+            current_lr = optimizer.param_groups[0]['lr']
+            writer.add_scalar('Learning_Rate', current_lr, global_step + i)
+
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        if i % args.print_freq == 0:
+            progress.display(i + 1)
+
+        # Perform multiple verifications within an epoch
+        if (i + 1) % validate_freq == 0:
+            model.eval()
+
+            results = validate(val_loader=val_loader, model=model, criterion=criterion, args=args, logger=logger)
+            val_srcc = results['metrics']['srcc']
+            val_plcc = results['metrics']['plcc']
+            val_klcc = results['metrics']['klcc']
+            val_loss = results['metrics']['loss']
+            logger.info(f'Validation results: SRCC: {val_srcc:.4f}, PLCC: {val_plcc:.4f}, KLCC: {val_klcc:.4f}, Loss: {val_loss:.4f}')
+            if args.rank == 0:
+                # Add the validation loss to the same loss chart
+                writer.add_scalars('Loss', {
+                    'val': val_loss
+                }, global_step + i)
+
+                # Add all performance metrics to the same Metrics chart.
+                writer.add_scalars('Metrics', {
+                    'SRCC': val_srcc,
+                    'PLCC': val_plcc,
+                    'KLCC': val_klcc
+                }, global_step + i)
+
+            is_best = val_srcc > best_srcc
+            best_srcc = max(val_srcc, best_srcc)
+            best_plcc = max(val_plcc, best_plcc)
+            best_klcc = max(val_klcc, best_klcc)
+
+            # Save the best model and results
+            if not args.multiprocessing_distributed or (args.multiprocessing_distributed
+                                                        and args.rank % args.ngpus_per_node == 0):
+                save_checkpoint(args, {
+                    'epoch': epoch + 1,
+                    'arch': args.arch,
+                    'state_dict': model.state_dict(),
+                    'best_srcc': best_srcc,
+                    'optimizer': optimizer.state_dict(),
+                    'scheduler': scheduler.state_dict()
+                }, is_best)
+
+                if is_best:
+                    logger.info(
+                        f'**BEST** Validation results: SRCC: {best_srcc:.4f}, PLCC: {best_plcc:.4f}, KLCC: {best_klcc:.4f}')
+
+                    df = pd.DataFrame({
+                        'image_name': results['image_names'],
+                        'prediction': results['predictions'],
+                        'ground_truth': results['ground_truth']
+                    })
+                    csv_filename = os.path.join(args.output_dir,
+                                                f'{args.run_name}_best_val_results.csv')
+                    df.to_csv(csv_filename, index=False)
+
+            model.train()
+
+    logger.info(
+        f'**BEST** Validation results: SRCC: {best_srcc:.4f}, PLCC: {best_plcc:.4f}, KLCC: {best_klcc:.4f}')
+    return best_srcc, best_plcc, best_klcc
+
+@torch.no_grad()
+def validate(val_loader, model, args, criterion, logger):
+    model.eval()
+    val_dataset_len = len(val_loader.dataset)
+    val_loader_len = len(val_loader)
+    batch_time = AverageMeter('Time', ':6.3f')
+    losses = AverageMeter('Loss', ':.4e')
+
+    with torch.no_grad():
+        temp_pred_scores = []
+        temp_gt_scores = []
+        temp_img_names = []
+        time_tmp = time.time()
+
+        for i, batch in enumerate(val_loader):
+
+            if args.gpu is not None and torch.cuda.is_available():
+                device = torch.device(f'cuda:{args.gpu}')
+            elif torch.backends.mps.is_available():
+                device = torch.device('mps')
+            else:
+                device = torch.device('cpu')
+
+            images = batch['image_cropped'].to(device, non_blocking=True if device.type == 'cuda' else False)
+            target = batch['label'].to(device, non_blocking=True if device.type == 'cuda' else False)
+
+            if 'image_resized' in batch:
+                image_resized = batch['image_resized'].to(device, non_blocking=True if device.type == 'cuda' else False)
+                output = model(images, image_resized).view(-1)
+            else:
+                output = model(images).view(-1)
+
+            # if args.gpu is not None and torch.cuda.is_available():
+            #     images = batch['image'].cuda(args.gpu, non_blocking=True)
+            #     target = batch['label'].cuda(args.gpu, non_blocking=True)
+            # if torch.backends.mps.is_available():
+            #     images = images.to('mps')
+            #     target = target.to('mps')
+
+            # output = model(images).view(-1)
+            loss = criterion(output, target.view(-1))
+            loss = reduce_tensor(loss)
+            losses.update(loss.item(), target.size(0))
+
+            batch_time.update(time.time() - time_tmp)
+            time_tmp = time.time()
+
+            # Save predicted values, gt values, and image names
+            temp_pred_scores.append(output.view(-1))
+            temp_gt_scores.append(target.view(-1))
+            temp_img_names.extend(batch['image_name'])
+
+            if i % args.print_freq == 0:
+                logger.info(
+                    f"Test: [{i}/{val_loader_len}]\t"
+                    f"Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t"
+                    f"Loss {losses.val:.4f} ({losses.avg:.4f})\t"
+                )
+
+    # Combine the results of all batches
+    pred_scores = torch.cat(temp_pred_scores)
+    gt_scores = torch.cat(temp_gt_scores)
+
+    # Distributed processing
+    if torch.distributed.is_initialized():
+        # Collect the results of all processes
+        img_names_gather = [None for _ in range(dist.get_world_size())]
+        torch.distributed.all_gather_object(img_names_gather, temp_img_names)
+        all_img_names = []
+        for names in img_names_gather:
+            all_img_names.extend(names)
+        all_img_names = all_img_names[:val_dataset_len]  # 截取到实际数据集大小
+
+        preds_gather_list = [
+            torch.zeros_like(pred_scores) for _ in range(dist.get_world_size())
+        ]
+        torch.distributed.all_gather(preds_gather_list, pred_scores)
+        gather_preds = torch.cat(preds_gather_list, dim=0)[:val_dataset_len]
+
+        grotruth_gather_list = [
+            torch.zeros_like(gt_scores) for _ in range(dist.get_world_size())
+        ]
+        torch.distributed.all_gather(grotruth_gather_list, gt_scores)
+        gather_grotruth = torch.cat(grotruth_gather_list, dim=0)[:val_dataset_len]
+
+        if args.patch_num > 1:
+            gather_preds_matrix = gather_preds.view(-1, args.patch_num)
+
+            gather_preds = (gather_preds_matrix.mean(dim=-1)).squeeze()
+            gather_grotruth = (
+                (gather_grotruth.view(-1, args.patch_num)).mean(dim=-1)
+            ).squeeze()
+
+        final_preds = gather_preds.float().detach()
+        final_grotruth = gather_grotruth.float().detach()
+    else:
+        final_preds = pred_scores.float().detach()
+        final_grotruth = gt_scores.float().detach()
+        all_img_names = temp_img_names
+
+    # Calculate the correlation coefficient
+    try:
+        logger.info(f"len of dataset: {val_dataset_len}, final_preds shape: {final_preds.shape}, final_grotruth shape: {final_grotruth.shape}")
+        # Check for the presence of NaN or inf
+        if torch.isnan(final_preds).any() or torch.isinf(final_preds).any() or \
+                torch.isnan(final_grotruth).any() or torch.isinf(final_grotruth).any():
+            raise ValueError("Found NaN or inf values in predictions or ground truth")
+
+        test_srcc = torchmetrics.functional.spearman_corrcoef(final_preds, final_grotruth).item()
+        test_plcc = torchmetrics.functional.pearson_corrcoef(final_preds, final_grotruth).item()
+        test_klcc = torchmetrics.functional.kendall_rank_corrcoef(final_preds, final_grotruth).item()
+
+    except Exception as e:
+        logger.warning(f"Error in calculating correlations: {str(e)}. Resetting cc relation to zero...")
+        test_plcc = 0.0
+        test_srcc = 0.0
+        test_klcc = 0.0
+
+    # Create a result dictionary containing the correspondence between image names, predicted values, and actual values.
+    results = {
+        'image_names': all_img_names,
+        'predictions': final_preds.cpu().numpy().tolist(),
+        'ground_truth': final_grotruth.cpu().numpy().tolist(),
+        'metrics': {
+            'srcc': test_srcc,
+            'plcc': test_plcc,
+            'klcc': test_klcc,
+            'loss': losses.avg
+        }
+    }
+
+    return results
+
+class Summary(Enum):
+    NONE = 0
+    AVERAGE = 1
+    SUM = 2
+    COUNT = 3
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self, name, fmt=':f', summary_type=Summary.AVERAGE):
+        self.name = name
+        self.fmt = fmt
+        self.summary_type = summary_type
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+    def all_reduce(self):
+        if torch.cuda.is_available():
+            device = torch.device("cuda")
+        elif torch.backends.mps.is_available():
+            device = torch.device("mps")
+        else:
+            device = torch.device("cpu")
+        total = torch.tensor([self.sum, self.count], dtype=torch.float32, device=device)
+        dist.all_reduce(total, dist.ReduceOp.SUM, async_op=False)
+        self.sum, self.count = total.tolist()
+        self.avg = self.sum / self.count
+
+    def __str__(self):
+        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
+        return fmtstr.format(**self.__dict__)
+
+    def summary(self):
+        fmtstr = ''
+        if self.summary_type is Summary.NONE:
+            fmtstr = ''
+        elif self.summary_type is Summary.AVERAGE:
+            fmtstr = '{name} {avg:.3f}'
+        elif self.summary_type is Summary.SUM:
+            fmtstr = '{name} {sum:.3f}'
+        elif self.summary_type is Summary.COUNT:
+            fmtstr = '{name} {count:.3f}'
+        else:
+            raise ValueError('invalid summary type %r' % self.summary_type)
+
+        return fmtstr.format(**self.__dict__)
+
+
+class ProgressMeter(object):
+    def __init__(self, num_batches, meters, prefix=""):
+        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
+        self.meters = meters
+        self.prefix = prefix
+
+    def display(self, batch):
+        entries = [self.prefix + self.batch_fmtstr.format(batch)]
+        entries += [str(meter) for meter in self.meters]
+        print('\t'.join(entries))
+
+    def display_summary(self):
+        entries = [" *"]
+        entries += [meter.summary() for meter in self.meters]
+        print(' '.join(entries))
+
+    def _get_batch_fmtstr(self, num_batches):
+        num_digits = len(str(num_batches // 1))
+        fmt = '{:' + str(num_digits) + 'd}'
+        return '[' + fmt + '/' + fmt.format(num_batches) + ']'
+
+
+def accuracy(output, target, topk=(1,)):
+    """Computes the accuracy over the k top predictions for the specified values of k"""
+    with torch.no_grad():
+        maxk = max(topk)
+        batch_size = target.size(0)
+
+        _, pred = output.topk(maxk, 1, True, True)
+        pred = pred.t()
+        correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+        res = []
+        for k in topk:
+            correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
+            res.append(correct_k.mul_(100.0 / batch_size))
+        return res
+
+
+if __name__ == '__main__':
+
+    args = get_args().parse_args()
+
+    args.run_name = args.arch + '_' + args.dataset + '_' + args.metric_type
+
+    os.makedirs(os.path.join(args.output_dir), exist_ok=True)
+    os.makedirs(os.path.join(args.output_dir, 'tensorboard_logs', args.run_name), exist_ok=True)
+
+    # save config file
+    with open(os.path.join(args.output_dir, 'config.yaml'), 'w') as f:
+        f.write(args.__dict__.__str__())
+
+    main(args)
+
+
+
+