# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import time

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as TF
from tabulate import tabulate

import megengine as mge
import megengine.functional as MF
import megengine.module as MM

module_cache = {
    "conv2d": (MM.Conv2d(32, 32, 3, 1, 0), nn.Conv2d(32, 32, 3, 1, 0).cuda()),
    "dw_conv2d": (
        MM.Conv2d(32, 32, 3, 1, 0, groups=32),
        nn.Conv2d(32, 32, 3, 1, 0, groups=32).cuda(),
    ),
    "conv3d": (MM.Conv3d(32, 32, 3, 1, 0), nn.Conv3d(32, 32, 3, 1, 0).cuda()),
    "ConvTranspose2d": (
        MM.ConvTranspose2d(32, 32, 3, 1, 0),
        nn.ConvTranspose2d(32, 32, 3, 1, 0).cuda(),
    ),
    "BatchNorm2d": (
MM.BatchNorm2d(64)