Commit a20c741b by Tianqi Chen Committed by tqchen

pickle memoize no longer print message (#1111)

parent 47ed25e8
...@@ -5,6 +5,7 @@ import sys ...@@ -5,6 +5,7 @@ import sys
import atexit import atexit
from decorator import decorate from decorator import decorate
from .._ffi.base import string_types from .._ffi.base import string_types
try: try:
import cPickle as pickle import cPickle as pickle
except ImportError: except ImportError:
...@@ -79,7 +80,6 @@ def memoize(key): ...@@ -79,7 +80,6 @@ def memoize(key):
else: else:
assert isinstance(arg, allow_types) assert isinstance(arg, allow_types)
if key in cache.cache: if key in cache.cache:
print("Use memoize {0}{1}".format(fkey, key))
return cache.cache[key] return cache.cache[key]
res = func(*args) res = func(*args)
cache.cache[key] = res cache.cache[key] = res
......
...@@ -2,5 +2,5 @@ ...@@ -2,5 +2,5 @@
"""Raspberry pi specific declaration and schedules.""" """Raspberry pi specific declaration and schedules."""
from __future__ import absolute_import as _abs from __future__ import absolute_import as _abs
from .conv2d import * from .conv2d import schedule_conv2d_nchw
from .depthwise_conv2d import * from .depthwise_conv2d import schedule_depthwise_conv2d_nchw
...@@ -4,7 +4,7 @@ from __future__ import absolute_import as _abs ...@@ -4,7 +4,7 @@ from __future__ import absolute_import as _abs
import tvm import tvm
from tvm import target as _target from tvm import target as _target
from .. import tag from .. import tag
from ..nn.conv2d import conv2d, _get_schedule from ..nn.conv2d import conv2d as _conv2d, _get_schedule
from ..nn.conv2d import SpatialPack, Im2ColPack from ..nn.conv2d import SpatialPack, Im2ColPack
from ..nn.conv2d import _WORKLOADS, _SCH_TO_DECL_FUNC from ..nn.conv2d import _WORKLOADS, _SCH_TO_DECL_FUNC
from ..nn.conv2d import _get_workload from ..nn.conv2d import _get_workload
...@@ -66,7 +66,7 @@ _SCHEDULES = [ ...@@ -66,7 +66,7 @@ _SCHEDULES = [
] ]
@_get_schedule.register("rasp") @_get_schedule.register("rasp")
def _schedule_conv2d(wkl): def _get_schedule_conv2d(wkl):
if wkl not in _WORKLOADS: if wkl not in _WORKLOADS:
raise ValueError("no schedule for such workload: {}".format(wkl)) raise ValueError("no schedule for such workload: {}".format(wkl))
idx = _WORKLOADS.index(wkl) idx = _WORKLOADS.index(wkl)
...@@ -74,7 +74,7 @@ def _schedule_conv2d(wkl): ...@@ -74,7 +74,7 @@ def _schedule_conv2d(wkl):
return sch return sch
@conv2d.register("rasp") @_conv2d.register("rasp")
def _declaration_conv2d(data, kernel, stride, padding, layout, out_dtype): def _declaration_conv2d(data, kernel, stride, padding, layout, out_dtype):
if out_dtype is None: if out_dtype is None:
out_dtype = data.dtype out_dtype = data.dtype
...@@ -309,7 +309,7 @@ def _schedule_im2col_conv2d(s, data, data_pad, data_col, data_vec, ...@@ -309,7 +309,7 @@ def _schedule_im2col_conv2d(s, data, data_pad, data_col, data_vec,
return s return s
@generic.schedule_conv2d_nchw.register(["rasp"]) @generic.schedule_conv2d_nchw.register(["rasp"])
def schedule_conv2d(outs): def schedule_conv2d_nchw(outs):
"""Create schedule for tensors""" """Create schedule for tensors"""
s = tvm.create_schedule([x.op for x in outs]) s = tvm.create_schedule([x.op for x in outs])
......
...@@ -164,7 +164,7 @@ def _schedule(s, data, data_pad, kernel, output, last): ...@@ -164,7 +164,7 @@ def _schedule(s, data, data_pad, kernel, output, last):
@generic.schedule_depthwise_conv2d_nchw.register(["cpu", "rasp"]) @generic.schedule_depthwise_conv2d_nchw.register(["cpu", "rasp"])
def schedule_depthwise_conv2d(outs): def schedule_depthwise_conv2d_nchw(outs):
"""Schedule for depthwise_conv2d nchw forward. """Schedule for depthwise_conv2d nchw forward.
Parameters Parameters
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment