v1.0

docker/requirements.txt

+# torch
+timm==0.5.4 
+fvcore
+onnx
+wandb

docker_start.sh

+docker run -it --shm-size=32G -v $PWD/RepViT:/home/RepViT -v /parastor/DL_DATA/HOT:/home/HOT -v /opt/hyhal:/opt/hyhal:ro --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --name repvit c85ed27005f2 bash
+# python -m torch.utils.collect_env

docker_startnv.sh

+docker run -it --shm-size=100G -v $PWD/RepViT:/home/RepViT -v /parastor/DL_DATA/HOT:/home/HOT --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --name repvit f6b99c8a0f01 bash
+# python -m torch.utils.collect_env

engine.py

+"""
+Train and eval functions used in main.py
+"""
+import math
+import sys
+from typing import Iterable, Optional
+
+import torch
+
+from timm.data import Mixup
+from timm.utils import accuracy, ModelEma
+
+from losses import DistillationLoss
+import utils
+
+def set_bn_state(model):
+    for m in model.modules():
+        if isinstance(m, torch.nn.modules.batchnorm._BatchNorm):
+            m.eval()
+
+def train_one_epoch(model: torch.nn.Module, criterion: DistillationLoss,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler,
+                    clip_grad: float = 0,
+                    clip_mode: str = 'norm',
+                    model_ema: Optional[ModelEma] = None, mixup_fn: Optional[Mixup] = None,
+                    set_training_mode=True,
+                    set_bn_eval=False,):
+    model.train(set_training_mode)
+    if set_bn_eval:
+        set_bn_state(model)
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', utils.SmoothedValue(
+        window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 100
+
+    for samples, targets in metric_logger.log_every(
+            data_loader, print_freq, header):
+        samples = samples.to(device, non_blocking=True)
+        targets = targets.to(device, non_blocking=True)
+
+        if mixup_fn is not None:
+            samples, targets = mixup_fn(samples, targets)
+
+        with torch.cuda.amp.autocast():
+            outputs = model(samples)
+            loss = criterion(samples, outputs, targets)
+
+        loss_value = loss.item()
+
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        optimizer.zero_grad()
+
+        # this attribute is added by timm on one optimizer (adahessian)
+        is_second_order = hasattr(
+            optimizer, 'is_second_order') and optimizer.is_second_order
+        loss_scaler(loss, optimizer, clip_grad=clip_grad, clip_mode=clip_mode,
+                    parameters=model.parameters(), create_graph=is_second_order)
+
+        torch.cuda.synchronize()
+        if model_ema is not None:
+            model_ema.update(model)
+
+        metric_logger.update(loss=loss_value)
+        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+@torch.no_grad()
+def evaluate(data_loader, model, device):
+    criterion = torch.nn.CrossEntropyLoss()
+
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    header = 'Test:'
+
+    # switch to evaluation mode
+    model.eval()
+
+    for images, target in metric_logger.log_every(data_loader, 10, header):
+        images = images.to(device, non_blocking=True)
+        target = target.to(device, non_blocking=True)
+
+        # compute output
+        with torch.cuda.amp.autocast():
+            output = model(images)
+            loss = criterion(output, target)
+
+        acc1, acc5 = accuracy(output, target, topk=(1, 5))
+
+        batch_size = images.shape[0]
+        metric_logger.update(loss=loss.item())
+        metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
+        metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print('* Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f} loss {losses.global_avg:.3f}'
+          .format(top1=metric_logger.acc1, top5=metric_logger.acc5, losses=metric_logger.loss))
+
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}

eval.sh

+python main.py --eval --model repvit_m1_1 --resume pretrain/repvit_m1_1_distill_300e.pth --data-path ~/imagenet
\ No newline at end of file

export_coreml.py

+import torch
+
+from timm import create_model
+import model
+
+import utils
+
+import torch
+import torchvision
+from argparse import ArgumentParser
+
+parser = ArgumentParser()
+
+parser.add_argument('--model', default='repvit_m1_1', type=str)
+parser.add_argument('--resolution', default=224, type=int)
+parser.add_argument('--ckpt', default=None, type=str)
+
+if __name__ == "__main__":
+    # Load a pre-trained version of MobileNetV2
+    args = parser.parse_args()
+    model = create_model(args.model, distillation=True)
+    # model = create_model(args.model, distillation=False)
+    # print(torch.load(args.ckpt)['model'])
+    if args.ckpt:
+        model.load_state_dict(torch.load(args.ckpt)['model'])
+    utils.replace_batchnorm(model)
+    model.eval()
+
+    # Trace the model with random data.
+    resolution = args.resolution
+    example_input = torch.rand(1, 3, resolution, resolution)
+    traced_model = torch.jit.trace(model, example_input)
+    torch.onnx.export(traced_model, example_input, "model.onnx", verbose=True, input_names=["input"], output_names=["output"])
+    out = traced_model(example_input)
+    '''
+
+    import coremltools as ct
+
+    # Using image_input in the inputs parameter:
+    # Convert to Core ML neural network using the Unified Conversion API.
+    model = ct.convert(
+        traced_model,
+        inputs=[ct.ImageType(shape=example_input.shape)]
+    )
+
+    # Save the converted model.
+    model.save(f"coreml/{args.model}_{resolution}.mlmodel")
+    '''

exportonnx.sh

+python export_coreml.py --ckpt checkpoints/repvit_m0_9/2024_05_22_11_20_59/checkpoint_0.pth --model repvit_m0_9

flops.py

+import torch
+import time
+from timm import create_model
+import model
+import utils
+from fvcore.nn import FlopCountAnalysis
+
+T0 = 5
+T1 = 10
+
+for n, batch_size, resolution in [
+    ('repvit_m0_9', 1024, 224),
+]:
+    inputs = torch.randn(1, 3, resolution,
+                            resolution)
+    model = create_model(n, num_classes=1000)
+    utils.replace_batchnorm(model)
+    n_parameters = sum(p.numel()
+                       for p in model.parameters() if p.requires_grad)
+    print('number of params:', n_parameters / 1e6)
+    flops = FlopCountAnalysis(model, inputs)
+    print("flops: ", flops.total() / 1e9)
\ No newline at end of file