main_benchmark.py

#!/usr/bin/env python3
# Copyright (c) 2014-2022 Megvii Inc. All rights reserved.
"""Benchmark depth-wise large convolution for megengine and pytorch"""
import os
import time

import megengine
import megengine.functional as F
import torch

torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True

assert megengine.__version__ >= "1.8.2", "please update megengine via " \
"`pip3 megengine -f https://megengine.org.cn/whl/mge.html -U --user`"


def benchmark_megengine(batch_size, resolution, channels, depth, kernel_size, niters=10):
    input = F.ones([batch_size, channels, resolution, resolution]) * 1e-3
    weight = F.ones([channels, 1, 1, kernel_size, kernel_size]) * 1e-3
    diff = []
    for i in range(niters):
        x = input
        megengine._full_sync()
        t = time.perf_counter()
        for _ in range(depth):
            x = F.conv2d(x, weight, bias=None, padding=kernel_size // 2, groups=channels)
        megengine._full_sync()
        diff.append((time.perf_counter() - t) * 1000)
    diff = sum(sorted(diff)[1:-2]) / (niters - 3)
    print(f"benchmark_megeg\tB{batch_size},R{resolution},C{channels},D{depth},K{kernel_size}\t{diff:.3f} ms")
    return diff


@torch.no_grad()
def benchmark_torch(batch_size, resolution, channels, depth, kernel_size, niters=10):
    input = torch.randn(batch_size, channels, resolution, resolution).cuda() * 1e-3
    weight = torch.randn(channels, 1, kernel_size, kernel_size).cuda() * 1e-3
    diff = []
    for i in range(niters):
        x = input.clone()
        torch.cuda.synchronize()
        t = time.perf_counter()
        for _ in range(depth):
            x = torch.nn.functional.conv2d(x, weight, bias=None, padding=kernel_size // 2, groups=channels)
        torch.cuda.synchronize()
        diff.append((time.perf_counter() - t) * 1000)
    diff = sum(sorted(diff)[1:-2]) / (niters - 3)
    print(f"benchmark_torch\tB{batch_size},R{resolution},C{channels},D{depth},K{kernel_size}\t{diff:.3f} ms")
    return diff


if __name__ == "__main__":
    for resolution in (16, 32):
        for kernel_size in (7, 17, 27):
            benchmark_torch(64, resolution, 256, 12, kernel_size)
            benchmark_megengine(64, resolution, 256, 12, kernel_size)