Skip to content

Instantly share code, notes, and snippets.

@unnonouno
Created July 5, 2016 08:31
Show Gist options
  • Save unnonouno/9fbeef09bb2552242a692c7dcc23af6a to your computer and use it in GitHub Desktop.
Save unnonouno/9fbeef09bb2552242a692c7dcc23af6a to your computer and use it in GitHub Desktop.
Mixing CPU/GPU parameters in one Chainer model
import chainer
from chainer import cuda
import chainer.functions as F
import chainer.links as L
import chainer.optimizers as O
import cupy
import numpy
class MyModel(chainer.Chain):
def __init__(self):
super(MyModel, self).__init__(
embed=L.EmbedID(10, 3),
W=L.Linear(3, 3),
)
def to_gpu(self, device=None):
with cuda.get_device(device):
self.W.to_gpu() # only send `W` to GPU
def __call__(self, x, y):
e = self.embed(x)
e = F.copy(e, 0) # send `e` to GPU0
return F.sum(self.W(e) - y)
model = MyModel()
model.to_gpu()
optimizer = O.SGD()
optimizer.setup(model)
x = numpy.array([1, 2, 3], 'i')
y = cupy.arange(9, dtype='f').reshape((3, 3))
l = model(x, y)
model.zerograds()
l.backward()
optimizer.update()
print(model.embed.W.data)
print(type(model.embed.W.data))
print(model.W.W.data)
print(type(model.W.W.data))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment