init

mambavision/validate_pip.sh

+#!/bin/bash
+DATA_PATH="/ImageNet/val"
+BS=128
+
+python validate_pip_model.py --model mamba_vision_T --data_dir=$DATA_PATH --batch-size $BS
+

mambavision/validate_pip_model.py

+#!/usr/bin/env python3
+""" ImageNet Validation Script
+
+This is intended to be a lean and easily modifiable ImageNet validation script for evaluating pretrained
+models or training checkpoints against ImageNet or similarly organized image datasets. It prioritizes
+canonical PyTorch, standard Python style, and good performance. Repurpose as you see fit.
+
+Hacked together by Ross Wightman (https://github.com/rwightman)
+"""
+import argparse
+import csv
+import glob
+import json
+import logging
+import os
+import time
+from collections import OrderedDict
+from contextlib import suppress
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+from timm.data import create_dataset, create_loader, resolve_data_config, RealLabelsImagenet
+from timm.models import create_model, load_checkpoint, is_model, list_models
+from timm.utils import accuracy, AverageMeter, natural_key, setup_default_logging, set_jit_fuser, \
+    decay_batch_step, check_batch_size_retry, ParseKwargs
+
+try:
+    from apex import amp
+    has_apex = True
+except ImportError:
+    has_apex = False
+
+has_native_amp = False
+try:
+    if getattr(torch.cuda.amp, 'autocast') is not None:
+        has_native_amp = True
+except AttributeError:
+    pass
+
+try:
+    from functorch.compile import memory_efficient_fusion
+    has_functorch = True
+except ImportError as e:
+    has_functorch = False
+
+has_compile = hasattr(torch, 'compile')
+
+_logger = logging.getLogger('validate')
+
+
+parser = argparse.ArgumentParser(description='PyTorch ImageNet Validation')
+parser.add_argument('data', nargs='?', metavar='DIR', const=None,
+                    help='path to dataset (*deprecated*, use --data-dir)')
+parser.add_argument('--data-dir', metavar='DIR',
+                    help='path to dataset (root dir)')
+parser.add_argument('--dataset', metavar='NAME', default='',
+                    help='dataset type + name ("<type>/<name>") (default: ImageFolder or ImageTar if empty)')
+parser.add_argument('--split', metavar='NAME', default='validation',
+                    help='dataset split (default: validation)')
+parser.add_argument('--dataset-download', action='store_true', default=False,
+                    help='Allow download of dataset for torch/ and tfds/ datasets that support it.')
+parser.add_argument('--model', '-m', metavar='NAME', default='dpn92',
+                    help='model architecture (default: dpn92)')
+parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
+                    help='number of data loading workers (default: 4)')
+parser.add_argument('-b', '--batch-size', default=256, type=int,
+                    metavar='N', help='mini-batch size (default: 256)')
+parser.add_argument('--img-size', default=None, type=int,
+                    metavar='N', help='Input image dimension, uses model default if empty')
+parser.add_argument('--in-chans', type=int, default=None, metavar='N',
+                    help='Image input channels (default: None => 3)')
+parser.add_argument('--input-size', default=None, nargs=3, type=int,
+                    metavar='N N N', help='Input all image dimensions (d h w, e.g. --input-size 3 224 224), uses model default if empty')
+parser.add_argument('--use-train-size', action='store_true', default=False,
+                    help='force use of train input size, even when test size is specified in pretrained cfg')
+parser.add_argument('--crop-pct', default=None, type=float,
+                    metavar='N', help='Input image center crop pct')
+parser.add_argument('--crop-mode', default=None, type=str,
+                    metavar='N', help='Input image crop mode (squash, border, center). Model default if None.')
+parser.add_argument('--mean', type=float, nargs='+', default=None, metavar='MEAN',
+                    help='Override mean pixel value of dataset')
+parser.add_argument('--std', type=float,  nargs='+', default=None, metavar='STD',
+                    help='Override std deviation of of dataset')
+parser.add_argument('--interpolation', default='', type=str, metavar='NAME',
+                    help='Image resize interpolation type (overrides model)')
+parser.add_argument('--num-classes', type=int, default=None,
+                    help='Number classes in dataset')
+parser.add_argument('--class-map', default='', type=str, metavar='FILENAME',
+                    help='path to class to idx mapping file (default: "")')
+parser.add_argument('--gp', default=None, type=str, metavar='POOL',
+                    help='Global pool type, one of (fast, avg, max, avgmax, avgmaxc). Model default if None.')
+parser.add_argument('--log-freq', default=10, type=int,
+                    metavar='N', help='batch logging frequency (default: 10)')
+parser.add_argument('--checkpoint', default='', type=str, metavar='PATH',
+                    help='path to latest checkpoint (default: none)')
+parser.add_argument('--pretrained', dest='pretrained', action='store_true',
+                    help='use pre-trained model')
+parser.add_argument('--num-gpu', type=int, default=1,
+                    help='Number of GPUS to use')
+parser.add_argument('--test-pool', dest='test_pool', action='store_true',
+                    help='enable test time pool')
+parser.add_argument('--no-prefetcher', action='store_true', default=False,
+                    help='disable fast prefetcher')
+parser.add_argument('--pin-mem', action='store_true', default=False,
+                    help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
+parser.add_argument('--channels-last', action='store_true', default=False,
+                    help='Use channels_last memory layout')
+parser.add_argument('--device', default='cuda', type=str,
+                    help="Device (accelerator) to use.")
+parser.add_argument('--amp', action='store_true', default=False,
+                    help='use NVIDIA Apex AMP or Native AMP for mixed precision training')
+parser.add_argument('--amp-dtype', default='float16', type=str,
+                    help='lower precision AMP dtype (default: float16)')
+parser.add_argument('--amp-impl', default='native', type=str,
+                    help='AMP impl to use, "native" or "apex" (default: native)')
+parser.add_argument('--tf-preprocessing', action='store_true', default=False,
+                    help='Use Tensorflow preprocessing pipeline (require CPU TF installed')
+parser.add_argument('--use-ema', dest='use_ema', action='store_true',
+                    help='use ema version of weights if present')
+parser.add_argument('--fuser', default='', type=str,
+                    help="Select jit fuser. One of ('', 'te', 'old', 'nvfuser')")
+parser.add_argument('--fast-norm', default=False, action='store_true',
+                    help='enable experimental fast-norm')
+parser.add_argument('--model-kwargs', nargs='*', default={}, action=ParseKwargs)
+
+
+scripting_group = parser.add_mutually_exclusive_group()
+scripting_group.add_argument('--torchscript', default=False, action='store_true',
+                             help='torch.jit.script the full model')
+scripting_group.add_argument('--torchcompile', nargs='?', type=str, default=None, const='inductor',
+                             help="Enable compilation w/ specified backend (default: inductor).")
+scripting_group.add_argument('--aot-autograd', default=False, action='store_true',
+                             help="Enable AOT Autograd support.")
+
+parser.add_argument('--results-file', default='', type=str, metavar='FILENAME',
+                    help='Output csv file for validation results (summary)')
+parser.add_argument('--results-format', default='csv', type=str,
+                    help='Format for results file one of (csv, json) (default: csv).')
+parser.add_argument('--real-labels', default='', type=str, metavar='FILENAME',
+                    help='Real labels JSON file for imagenet evaluation')
+parser.add_argument('--valid-labels', default='', type=str, metavar='FILENAME',
+                    help='Valid label indices txt file for validation of partial label space')
+parser.add_argument('--retry', default=False, action='store_true',
+                    help='Enable batch size decay & retry for single model validation')
+
+
+def validate(args):
+    # might as well try to validate something
+    args.pretrained = args.pretrained or not args.checkpoint
+    args.prefetcher = not args.no_prefetcher
+
+    if torch.cuda.is_available():
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.benchmark = True
+
+    device = torch.device(args.device)
+
+    # resolve AMP arguments based on PyTorch / Apex availability
+    use_amp = None
+    amp_autocast = suppress
+    if args.amp:
+        if args.amp_impl == 'apex':
+            assert has_apex, 'AMP impl specified as APEX but APEX is not installed.'
+            assert args.amp_dtype == 'float16'
+            use_amp = 'apex'
+            _logger.info('Validating in mixed precision with NVIDIA APEX AMP.')
+        else:
+            assert has_native_amp, 'Please update PyTorch to a version with native AMP (or use APEX).'
+            assert args.amp_dtype in ('float16', 'bfloat16')
+            use_amp = 'native'
+            amp_dtype = torch.bfloat16 if args.amp_dtype == 'bfloat16' else torch.float16
+            amp_autocast = partial(torch.autocast, device_type=device.type, dtype=amp_dtype)
+            _logger.info('Validating in mixed precision with native PyTorch AMP.')
+    else:
+        _logger.info('Validating in float32. AMP not enabled.')
+
+    if args.fuser:
+        set_jit_fuser(args.fuser)
+
+    if args.fast_norm:
+        set_fast_norm()
+
+    # create model
+    in_chans = 3
+    if args.in_chans is not None:
+        in_chans = args.in_chans
+    elif args.input_size is not None:
+        in_chans = args.input_size[0]
+
+    from mambavision import create_model
+    model = create_model(args.model, 
+                         pretrained=True,
+                         model_path="/tmp/model.pth.tar")
+    
+    if args.num_classes is None:
+        assert hasattr(model, 'num_classes'), 'Model must have `num_classes` attr if not set on cmd line/config.'
+        args.num_classes = model.num_classes
+
+    if args.checkpoint:
+        load_checkpoint(model, args.checkpoint, args.use_ema)
+
+    param_count = sum([m.numel() for m in model.parameters()])
+    _logger.info('Model %s created, param count: %d' % (args.model, param_count))
+
+    data_config = resolve_data_config(
+        vars(args),
+        model=model,
+        use_test_size=not args.use_train_size,
+        verbose=True,
+    )
+
+    test_time_pool = False
+    if args.test_pool:
+        model, test_time_pool = apply_test_time_pool(model, data_config)
+
+    model = model.to(device)
+    if args.channels_last:
+        model = model.to(memory_format=torch.channels_last)
+
+    if args.torchscript:
+        assert not use_amp == 'apex', 'Cannot use APEX AMP with torchscripted model'
+        model = torch.jit.script(model)
+    elif args.torchcompile:
+        assert has_compile, 'A version of torch w/ torch.compile() is required for --compile, possibly a nightly.'
+        torch._dynamo.reset()
+        model = torch.compile(model, backend=args.torchcompile)
+    elif args.aot_autograd:
+        assert has_functorch, "functorch is needed for --aot-autograd"
+        model = memory_efficient_fusion(model)
+
+    if use_amp == 'apex':
+        model = amp.initialize(model, opt_level='O1')
+
+    if args.num_gpu > 1:
+        model = torch.nn.DataParallel(model, device_ids=list(range(args.num_gpu)))
+
+    criterion = nn.CrossEntropyLoss().to(device)
+    DATA_PATH="/home/ali/Desktop/data_local/ImageNet-Validation/val"
+
+    root_dir = args.data or args.data_dir or DATA_PATH
+    dataset = create_dataset(
+        root=root_dir,
+        name=args.dataset,
+        split=args.split,
+        download=args.dataset_download,
+        load_bytes=args.tf_preprocessing,
+        class_map=args.class_map,
+    )
+
+    if args.valid_labels:
+        with open(args.valid_labels, 'r') as f:
+            valid_labels = [int(line.rstrip()) for line in f]
+    else:
+        valid_labels = None
+
+    if args.real_labels:
+        real_labels = RealLabelsImagenet(dataset.filenames(basename=True), real_json=args.real_labels)
+    else:
+        real_labels = None
+
+    crop_pct = 1.0 if test_time_pool else data_config['crop_pct']
+    loader = create_loader(
+        dataset,
+        input_size=data_config['input_size'],
+        batch_size=args.batch_size,
+        use_prefetcher=args.prefetcher,
+        interpolation=data_config['interpolation'],
+        mean=data_config['mean'],
+        std=data_config['std'],
+        num_workers=args.workers,
+        crop_pct=crop_pct,
+        crop_mode=data_config['crop_mode'],
+        pin_memory=args.pin_mem,
+        device=device,
+        tf_preprocessing=args.tf_preprocessing,
+    )
+
+    batch_time = AverageMeter()
+    losses = AverageMeter()
+    top1 = AverageMeter()
+    top5 = AverageMeter()
+
+    model.eval()
+    with torch.no_grad():
+        # warmup, reduce variability of first batch time, especially for comparing torchscript vs non
+        input = torch.randn((args.batch_size,) + tuple(data_config['input_size'])).to(device)
+        if args.channels_last:
+            input = input.contiguous(memory_format=torch.channels_last)
+        with amp_autocast():
+            model(input)
+
+        end = time.time()
+        for batch_idx, (input, target) in enumerate(loader):
+            if args.no_prefetcher:
+                target = target.to(device)
+                input = input.to(device)
+            if args.channels_last:
+                input = input.contiguous(memory_format=torch.channels_last)
+
+            # compute output
+            with amp_autocast():
+                output = model(input)
+
+                if valid_labels is not None:
+                    output = output[:, valid_labels]
+                loss = criterion(output, target)
+
+            if real_labels is not None:
+                real_labels.add_result(output)
+
+            # measure accuracy and record loss
+            acc1, acc5 = accuracy(output.detach(), target, topk=(1, 5))
+            losses.update(loss.item(), input.size(0))
+            top1.update(acc1.item(), input.size(0))
+            top5.update(acc5.item(), input.size(0))
+
+            # measure elapsed time
+            batch_time.update(time.time() - end)
+            end = time.time()
+
+            if batch_idx % args.log_freq == 0:
+                _logger.info(
+                    'Test: [{0:>4d}/{1}]  '
+                    'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s, {rate_avg:>7.2f}/s)  '
+                    'Loss: {loss.val:>7.4f} ({loss.avg:>6.4f})  '
+                    'Acc@1: {top1.val:>7.3f} ({top1.avg:>7.3f})  '
+                    'Acc@5: {top5.val:>7.3f} ({top5.avg:>7.3f})'.format(
+                        batch_idx,
+                        len(loader),
+                        batch_time=batch_time,
+                        rate_avg=input.size(0) / batch_time.avg,
+                        loss=losses,
+                        top1=top1,
+                        top5=top5
+                    )
+                )
+
+    if real_labels is not None:
+        # real labels mode replaces topk values at the end
+        top1a, top5a = real_labels.get_accuracy(k=1), real_labels.get_accuracy(k=5)
+    else:
+        top1a, top5a = top1.avg, top5.avg
+    results = OrderedDict(
+        model=args.model,
+        top1=round(top1a, 4), top1_err=round(100 - top1a, 4),
+        top5=round(top5a, 4), top5_err=round(100 - top5a, 4),
+        param_count=round(param_count / 1e6, 2),
+        img_size=data_config['input_size'][-1],
+        crop_pct=crop_pct,
+        interpolation=data_config['interpolation'],
+    )
+
+    _logger.info(' * Acc@1 {:.3f} ({:.3f}) Acc@5 {:.3f} ({:.3f})'.format(
+       results['top1'], results['top1_err'], results['top5'], results['top5_err']))
+
+    return results
+
+
+def _try_run(args, initial_batch_size):
+    batch_size = initial_batch_size
+    results = OrderedDict()
+    error_str = 'Unknown'
+    while batch_size:
+        args.batch_size = batch_size * args.num_gpu  # multiply by num-gpu for DataParallel case
+        try:
+            if torch.cuda.is_available() and 'cuda' in args.device:
+                torch.cuda.empty_cache()
+            results = validate(args)
+            return results
+        except RuntimeError as e:
+            error_str = str(e)
+            _logger.error(f'"{error_str}" while running validation.')
+            if not check_batch_size_retry(error_str):
+                break
+        batch_size = decay_batch_step(batch_size)
+        _logger.warning(f'Reducing batch size to {batch_size} for retry.')
+    results['error'] = error_str
+    _logger.error(f'{args.model} failed to validate ({error_str}).')
+    return results
+
+
+_NON_IN1K_FILTERS = ['*_in21k', '*_in22k', '*in12k', '*_dino', '*fcmae', '*seer']
+
+
+def main():
+    setup_default_logging()
+    args = parser.parse_args()
+    model_cfgs = []
+    model_names = []
+    if os.path.isdir(args.checkpoint):
+        # validate all checkpoints in a path with same model
+        checkpoints = glob.glob(args.checkpoint + '/*.pth.tar')
+        checkpoints += glob.glob(args.checkpoint + '/*.pth')
+        model_names = list_models(args.model)
+        model_cfgs = [(args.model, c) for c in sorted(checkpoints, key=natural_key)]
+    else:
+        if args.model == 'all':
+            # validate all models in a list of names with pretrained checkpoints
+            args.pretrained = True
+            model_names = list_models(
+                pretrained=True,
+                exclude_filters=_NON_IN1K_FILTERS,
+            )
+            model_cfgs = [(n, '') for n in model_names]
+        elif not is_model(args.model):
+            # model name doesn't exist, try as wildcard filter
+            model_names = list_models(
+                args.model,
+                pretrained=True,
+            )
+            model_cfgs = [(n, '') for n in model_names]
+
+        if not model_cfgs and os.path.isfile(args.model):
+            with open(args.model) as f:
+                model_names = [line.rstrip() for line in f]
+            model_cfgs = [(n, None) for n in model_names if n]
+
+    if len(model_cfgs):
+        _logger.info('Running bulk validation on these pretrained models: {}'.format(', '.join(model_names)))
+        results = []
+        try:
+            initial_batch_size = args.batch_size
+            for m, c in model_cfgs:
+                args.model = m
+                args.checkpoint = c
+                r = _try_run(args, initial_batch_size)
+                if 'error' in r:
+                    continue
+                if args.checkpoint:
+                    r['checkpoint'] = args.checkpoint
+                results.append(r)
+        except KeyboardInterrupt as e:
+            pass
+        results = sorted(results, key=lambda x: x['top1'], reverse=True)
+    else:
+        if args.retry:
+            results = _try_run(args, args.batch_size)
+        else:
+            results = validate(args)
+
+    if args.results_file:
+        write_results(args.results_file, results, format=args.results_format)
+
+    # output results in JSON to stdout w/ delimiter for runner script
+    print(f'--result\n{json.dumps(results, indent=4)}')
+
+
+def write_results(results_file, results, format='csv'):
+    with open(results_file, mode='w') as cf:
+        if format == 'json':
+            json.dump(results, cf, indent=4)
+        else:
+            if not isinstance(results, (list, tuple)):
+                results = [results]
+            if not results:
+                return
+            dw = csv.DictWriter(cf, fieldnames=results[0].keys())
+            dw.writeheader()
+            for r in results:
+                dw.writerow(r)
+            cf.flush()
+
+
+
+if __name__ == '__main__':
+    main()
\ No newline at end of file

model.properties

+# 模型唯一标识
+modelCode=1019
+# 模型名称
+modelName=MambaVision_pytorch
+# 模型描述
+modelDescription=MambaVision：混合 Mamba-Transformer 视觉主干。
+# 应用场景
+appScenario=推理,训练,图像分类,科研,教育,政府,金融
+# 框架类型
+frameType=Pytorch

requirements.txt

+torch>=2.4.0
+causal-conv1d
+mamba-ssm
+timm
+tensorboardX
+einops
+transformers

setup.cfg

+[dist_conda]
+
+conda_name_differences = 'torch:pytorch'
+channels = pytorch
+noarch = True
\ No newline at end of file

setup.py

+""" Setup
+"""
+from setuptools import setup, find_packages
+from codecs import open
+from os import path
+
+here = path.abspath(path.dirname(__file__))
+
+# Get the long description from the README file
+with open(path.join(here, 'README.md'), encoding='utf-8') as f:
+    long_description = f.read()
+
+setup(
+    name='mambavision',
+    version='1.0.9',
+    description='MambaVision: A Hybrid Mamba-Transformer Vision Backbone',
+    long_description=long_description,
+    long_description_content_type='text/markdown',
+    url='https://github.com/NVlabs/MambaVision',
+    author='Ali Hatamizadeh',
+    author_email='ahatamiz123@gmail.com',
+    classifiers=[
+        'Programming Language :: Python :: 3.7',
+        'Programming Language :: Python :: 3.8',
+        'Programming Language :: Python :: 3.9',
+        'Programming Language :: Python :: 3.10',
+        'Programming Language :: Python :: 3.11',
+        'Programming Language :: Python :: 3.12',
+        'Topic :: Scientific/Engineering :: Artificial Intelligence',
+        'Topic :: Software Development',
+        'Topic :: Software Development :: Libraries',
+        'Topic :: Software Development :: Libraries :: Python Modules',
+    ],
+
+    # Note that this is a string of words separated by whitespace, not a list.
+    keywords='pytorch pretrained models mamba vision transformer vit',
+    packages=find_packages(exclude=['tests']),
+    include_package_data=True,
+    install_requires=['torch >= 2.1.2', 'timm==1.0.9', 'mamba-ssm==2.2.2', 'causal-conv1d==1.4.0', 'einops'],
+    license="NVIDIA Source Code License-NC",
+    python_requires='>=3.7',
+)