inspect

#!/usr/bin/env python3

import sys
import torch
import math
from torch.utils.data import DataLoader

import torch.nn.functional as F

from resnet import *
from swiftnet import *
from karesnet import *

target='DPUCZDX8G_ISA1_B4096'
device = torch.device('cuda')

if len(sys.argv) == 1:
    command = 'inspect'
    device = torch.device('cpu')
else:
    command = sys.argv[1]

backbone = 'KAResNet9_no_square'
dummy_input = torch.randn([1, 3, 256, 512]).to(device)
dataset_root = f'../{dummy_input.shape[3]}x{dummy_input.shape[2]}cityscapes'

model_path = '_'.join([
    'swiftnet',
    backbone,
    dataset_root[3:],
    'nearest',
    '128',
    '19',
    'best',
    'model'
]) + '.pt'

model_path = "swiftnet_no_square_g2_k1_KAResNet9_512x256cityscapes_nearest_128_19_best_model.pt"

model = torch.load(model_path, map_location=torch.device('cpu'))

class ReLUKANFused(nn.Module):
    def __init__(self, module):
        super().__init__()

        self.base_conv = module.base_conv
        self.layer_norm = module.layer_norm
        self.relukan_conv = module.relukan_conv

        ch = module.phase_low.shape[1]
        self.sub_low = nn.Conv2d(ch, ch*(module.g+module.k), 1, groups=ch, bias=True)
        self.sub_high = nn.Conv2d(ch, ch*(module.g+module.k), 1, groups=ch, bias=True)

        sqrt_r = math.sqrt(module.r)
        self.sub_low.weight.data.fill_(sqrt_r)
        self.sub_high.weight.data.fill_(-sqrt_r)

        self.sub_low.bias.data = -module.phase_low.flatten() * sqrt_r
        self.sub_high.bias.data = module.phase_high.flatten() * sqrt_r

    def forward(self, x):
        basis = F.relu(self.base_conv(x), inplace=True)

        x1 = F.relu(self.sub_low(x), inplace=True).permute(0,2,3,1)
        x2 = F.relu(self.sub_high(x), inplace=True).permute(0,2,3,1)

        # x3 = x1 * x2
        x = x1 * x2
        # x = x * x3
        x = x.permute(0, 3, 1, 2)

        y = self.relukan_conv(x)
        y = F.relu(self.layer_norm(y + basis), inplace=True)
        return y

if 'KAResNet9' in backbone:
    model.backbone.conv1 = ReLUKANFused(model.backbone.conv1)
    model.backbone.layer1.conv1 = ReLUKANFused(model.backbone.layer1.conv1)
    model.backbone.layer2.conv1 = ReLUKANFused(model.backbone.layer2.conv1)
    model.backbone.layer3.conv1 = ReLUKANFused(model.backbone.layer3.conv1)
    model.backbone.layer4.conv1 = ReLUKANFused(model.backbone.layer4.conv1)

model = model.to(device)
model.eval()
print(sum(p.numel() for p in model.parameters()))

if str(device) == 'cuda':
    def dev_sync():
        torch.cuda.synchronize()
else:
    def dev_sync():
        pass

if command == 'inspect':
    from pytorch_nndct import Inspector
    inspector=Inspector(target)
    inspector.inspect(model, (dummy_input,), device=device, output_dir=backbone+'_inspect', image_format='png')
elif command == 'calib':
    from cityscapes import CityScapes
    from pytorch_nndct.apis import torch_quantizer
    dataset = CityScapes(dataset_root, 'val')
    loader = DataLoader(dataset, batch_size=1, shuffle=True, num_workers=4)

    quantizer = torch_quantizer('calib', model, dummy_input, output_dir=backbone+'_quantize', device=device)
    quant_model = quantizer.quant_model
    count = 0
    for i, _ in loader:
        quant_model(i.to(device))
        count += 1
        print(f'\r{count:3d}/{len(loader)}', end='')
    print('\n')
    quantizer.export_quant_config()
elif command == 'test':
    from cityscapes import CityScapes
    from pytorch_nndct.apis import torch_quantizer
    dataset = CityScapes(dataset_root, 'val')
    loader = DataLoader(dataset, batch_size=1, shuffle=True, num_workers=4)

    quantizer = torch_quantizer('test', model, dummy_input, output_dir=backbone+'_quantize', device=device)
    quantizer.quant_model(dummy_input)
    quantizer.export_xmodel(output_dir=backbone+'_quantize')
elif command == 'compile':
    import os
    os.system(' '.join([
        'vai_c_xir',
        '-x',
        backbone+'_quantize/SwiftNet_int.xmodel',
        '-a',
        '/opt/vitis_ai/compiler/arch/DPUCZDX8G/KV260/arch.json',
        '-o', backbone+'_compile',
    ]))
elif command == 'vis':
  from cityscapes import CityScapes
  c_inputs = []
  c_labels = []
  c_preds = []

  test_data = CityScapes(dataset_root, 'test')
  test_loader = DataLoader(test_data, batch_size=1, num_workers=2)

  for i, l in test_loader:
    i, l = i.to(device, dtype=torch.float32), l.to(device)
    with torch.no_grad():
      o = model(i)
    # e = evaluate(o.cpu(), l.cpu())

    for j in range(o.shape[0]):
      c_preds.append(o[j].cpu())
      c_inputs.append(i[j].cpu())
      c_labels.append(l[j].cpu())
    break

  import matplotlib.pyplot as plt

  def colorize(labels):
    x = CityScapes.train_id_to_id[labels]
    x = CityScapes.id_to_color[x]
    colors = np.array(colors)
    return colors[labels]

  def visualize(images, gt, preds, rows):
      fig, axs = plt.subplots(rows, 3, figsize=(10,20))

      axs[0].set_title('Original Image')
      axs[1].set_title('Overlay')
      axs[2].set_title('Predicted')

      for ax in axs.flat:
          ax.axis('off')

      print(rows)
      for i in range(rows):
        image_np = images[i].numpy().transpose(1, 2, 0)

        gt_indices = gt[i].numpy()
        pred_indices = torch.argmax(preds[i], dim=0).numpy()

        colored_gts = CityScapes.train_id_to_color[gt_indices]
        colored_preds = [CityScapes.train_id_to_color[i] for i in pred_indices]

        axs[0].imshow(image_np)
        axs[1].imshow(image_np)
        axs[1].imshow(colored_preds, alpha=0.5)
        axs[2].imshow(colored_preds)

      plt.tight_layout()
      plt.savefig('asdf.png')
      print('ok')

  visualize(c_inputs, c_labels, c_preds, min(len(c_inputs), 10))
else:
    from cityscapes import CityScapes
    from torchmetrics import JaccardIndex
    from pytorch_nndct.apis import torch_quantizer
    from torch.profiler import profile, record_function, ProfilerActivity

    dataset = CityScapes(dataset_root, 'val')
    loader = DataLoader(dataset, batch_size=1, shuffle=True, num_workers=4)

    if command == 'quant':
        quantizer = torch_quantizer('test', model, dummy_input, output_dir=backbone+'_quantize', device=device)
        model = quantizer.quant_model
    model.eval()

    metrics = JaccardIndex(
        task='multiclass',
        num_classes=CityScapes.num_classes,
        ignore_index=CityScapes.num_classes,
        average='none'
    ).to(device)

    with torch.no_grad():
        inp = dataset[0][0].unsqueeze(0).to(device)
        logits = model(inp)
        dev_sync()

        n = 0
        if str(device) == 'cuda':
            for i, l in loader:
                o = model(i.to(device))
                o = torch.argmax(o, dim=1)
                metrics.update(o, l.to(device))
                n += 1
                print(f'\r{n:3d}/{len(loader)}', end='')
            mIoU = metrics.compute().mean()
            print(f'\nmIoU: {mIoU}')
        dev_sync()

        # n = 0
        # if str(device) == 'cuda':
        #     activities = [ProfilerActivity.CUDA]
        # else:
        #     activities = [ProfilerActivity.CPU]
        #
        # with profile(activities=activities, record_shapes=False) as prof:
        #     with record_function('model_inference'):
        #         for inputs, _ in loader:
        #             inputs = inputs.to(device, dtype=torch.float)
        #             outputs = model(inputs)
        #             n += 1
        #             print(f'\r{n:3d}/{len(loader)}', end='')
        # if str(device) == 'cuda':
        #     print('\n', prof.key_averages().table(sort_by='cuda_time_total'))
        # else:
        #     print('\n', prof.key_averages().table(sort_by='cpu_time_total'))