提交 61f85841 authored 作者: Arnaud Bergeron's avatar Arnaud Bergeron

Flake8 fixes.

上级 63f8d7c2
import os.path
import numpy
import theano
from theano import Op, Apply, Variable, tensor
from theano import Apply, Variable, tensor
from theano.compile import optdb
from theano.compile.ops import shape_i
......@@ -9,7 +8,7 @@ from theano.gof import local_optimizer, COp
from theano.scalar import as_scalar, constant
from . import opt
from .basic_ops import (as_gpuarray_variable, gpu_alloc, gpu_from_host,
from .basic_ops import (as_gpuarray_variable, gpu_from_host,
host_from_gpu, GpuAllocEmpty)
from .opt_util import alpha_merge, output_merge
from .pycuda_helper import ensure_pycuda_context
......@@ -54,7 +53,7 @@ class Gemm16(COp):
COp.__init__(self, ["gemm16.c"], "gemm16")
self.relu = relu
# relu = True will require more work in optimizations.
assert self.relu == False
assert self.relu is False
self.inplace = inplace
if self.inplace:
self.destroy_map = {0: [0]}
......@@ -75,7 +74,7 @@ class Gemm16(COp):
return Apply(self, [C, alpha, A, B, beta], [C.type()])
def perform(self, node, inputs, outputs):
ctx = ensure_pycuda_context()
ensure_pycuda_context()
C, alpha, A, B, beta = inputs
# The nervana code does not support the case where both inputs
# are trans, so we need to copy one if them if that is the
......@@ -143,7 +142,7 @@ if (GpuKernel_init(&k_%(name)s, c->ops, c->ctx, 1, &bcode, &sz,
def c_init_code_struct(self, node, nodename, sub):
codel = [super(Gemm16, self).c_init_code_struct(node, nodename, sub)]
for name in self.KERN_NAMES:
codel.append("memset(&k_{0}, 0, sizeof(GpuKernel));".format(name));
codel.append("memset(&k_{0}, 0, sizeof(GpuKernel));".format(name))
codel.append("const char *bcode;")
codel.append("size_t sz;")
codel.append("PyGpuContextObject *c = pygpu_default_context();")
......@@ -175,6 +174,7 @@ def local_dot_to_gemm16(node):
shape_i(A, 0, fgraph), shape_i(B, 1, fgraph))
return [host_from_gpu(Gemm16()(C, 1.0, A, B, 0.0))]
@opt.register_opt()
@alpha_merge(Gemm16, alpha_in=1, beta_in=4, nd=2)
def local_gemm16_alpha_merge(node, *inputs):
......
......@@ -2,7 +2,6 @@ from functools import wraps
import numpy
import theano
from theano import scalar as scal, Constant
from theano.gof import local_optimizer
from theano.tensor import (DimShuffle, get_scalar_constant_value,
......@@ -13,11 +12,12 @@ from .elemwise import GpuDimShuffle, GpuElemwise
_one = scal.constant(numpy.asarray(1.0, dtype='float32'))
def grab_cpu_scalar(v, nd):
if v.owner is not None:
n = v.owner
if (isinstance(n.op, GpuDimShuffle) and
n.op.new_order == ('x',) * nd):
n.op.new_order == ('x',) * nd):
return host_from_gpu(n.inputs[0])
elif (isinstance(n.op, DimShuffle) and
n.op.new_order == ('x',) * nd):
......@@ -28,7 +28,7 @@ def grab_cpu_scalar(v, nd):
return None
else:
if (isinstance(v, Constant) and
v.broadcastable == (True,) * nd):
v.broadcastable == (True,) * nd):
return v.dimshuffle(())
......@@ -64,8 +64,8 @@ def alpha_merge(cls, alpha_in, beta_in, nd):
@wraps(maker)
def opt(node):
if (isinstance(node.op, GpuElemwise) and
node.op.scalar_op == scal.mul and
node.nin == 2):
node.op.scalar_op == scal.mul and
node.nin == 2):
targ = find_node(node.inputs[0], cls)
if targ is None:
targ = find_node(node.inputs[1], cls)
......@@ -88,8 +88,8 @@ def output_merge(cls, alpha_in, beta_in, out_in, nd):
@wraps(maker)
def opt(node):
if (isinstance(node.op, GpuElemwise) and
node.op.scalar_op == scal.add and
node.nin == 2):
node.op.scalar_op == scal.add and
node.nin == 2):
targ = find_node(node.inputs[0], cls)
W = node.inputs[1]
if targ is None:
......
......@@ -8,6 +8,7 @@ except ImportError:
pycuda_initialized = False
pycuda_context = None
def ensure_pycuda_context():
global pycuda_context, pycuda_initialized
if not pycuda_initialized:
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论