winograd_conv.py 文件源码

def __init__(self, op, lib, dtype,
                 N, C, K,
                 H, W, P, Q,
                 pad_h, pad_w, external=None, bprop=False):

        super(XpropWinograd_4x4_3x3, self).__init__(lib, dtype,
             N, C, K, 1, H, W, 1, 3, 3, 1, P, Q,
             0, pad_h, pad_w, 1,1,1, 1,1,1, bprop)

        SMs = _get_sm_count()

        self.autotune_key = " ".join(native_str(x) for x in (op + "_4x4_3x3",
           SMs, dtype.itemsize, N, C, K, H, W, P, Q))
        # insert Python version in filename to avoid Py2/Py3 incompatibilities in shelve
        self.autotune_db_file = os.path.join(lib.cache_dir, "autotune%d.db" % sys.version_info[0])

        # allow for .5 seconds worth of warmup when autotuning
        # assume 10 Tflops on 24 SMs
        self.warmup = min(max(int(5e12 / (P * Q * K * N * C * 9 * 2.0) * (SMs / 24.0)), 1), 1000)

        if external is None:
            self.init()
        else:
            # allow override for unit testing
            self.initialized = True
            self.init(autotune=1, external=external)

        lib.set_scratch_size(self.image_size, self.filter_trans.size, self.bsum.size)