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

Flake8 fixes.

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