Skip to content

Instantly share code, notes, and snippets.

@ajbrock
Created November 18, 2016 00:02
Show Gist options
  • Save ajbrock/d6914018d68e5f2adee6040824351e83 to your computer and use it in GitHub Desktop.
Save ajbrock/d6914018d68e5f2adee6040824351e83 to your computer and use it in GitHub Desktop.
class MyOp(theano.Op):
# Properties attribute
#itypes and otypes attributes are
#compulsory if make_node method is not defined.
#They're the type of input and output respectively
itypes = [cuda.CudaNdarrayType([False,False,False,False])]
otypes = [cuda.CudaNdarrayType([False]*4)]
def __init__(self,r):
self.r = r
super(MyOp,self).__init__()
def perform(self, node, inputs, output_storage):
input = inputs[0]
input_shape=input.shape
out = T.zeros((input_shape[0]//self.r**2,input_shape[1],self.r*input_shape[2],self.r*input_shape[3]))
batch_size = input.shape[0]//self.r**2
for x in xrange(self.r): # loop across all feature maps belonging to this channel
for y in xrange(self.r):
out=T.set_subtensor(out[:,:,x::self.r,y::self.r],input[batch_size*(self.r*x+y):batch_size*(self.r*x+y+1),:,:,:])
return out
# Other type of implementation
def infer_shape(self,node, input_shapes):
input_shape = input_shapes[0]
return [(input_shape[0]//self.r**2,input_shape[1],self.r*input_shape[2],self.r*input_shape[3])]
def grad(self, inputs, output_grads):
out = output_grads[0]
return T.concatenate([out[:,:,x::self.r,y::self.r] for x in xrange(self.r) for y in xrange(self.r)],axis=0)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment