I have created a model which inherits nn.Module
and I have transferred to model.cuda()
. While passing the parameters in forward
function i want one argument to be in cpu but since I have used model.cuda()
all arguments are in turn being converted to gpu i think. How do i force a tensor to live in cpu ?
Tensor.cpu()
will transfer to cpu but the point of forcing the tensor in cpu is because my tensor is a big matrix and transferring to gpu and then to cpu is not necessary.
You can partially choose cpu or gpu for each weight. Did you mean this?
model = models.resnet18()
model.cuda()
model.layer1[0].conv1.weight.type()
'torch.cuda.FloatTensor'
model.layer1[0].conv2.weight.type()
'torch.cuda.FloatTensor'
model.layer1[0].conv1.cpu()
Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
model.layer1[0].conv1.weight.type()
'torch.FloatTensor'
model.layer1[0].conv2.weight.type()
'torch.cuda.FloatTensor'
class Model(nn.Module):
def __init__(self):
pass
def forward(self, inputs1, inputs2, inputs3):
# inputs1, inputs2, inputs3 are already in gpu.
pass
model = Model()
model = torch.nn.DataParallel(model)
model = model.cuda()
Here the inputs1, inputs2, inputs3 are in gpu even though i pass them as just tensor with cpu because of the model.cuda()
. The inputs2 is so big that crashes the program with cuda out of memory
. Is there a way to force inputs2
to be in cpu without converting them to cpu
in the forward function. ?
def forward(self, *inputs, **kwargs):
if not self.device_ids:
return self.module(*inputs, **kwargs)
inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
if len(self.device_ids) == 1:
return self.module(*inputs[0], **kwargs[0])
replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
outputs = self.parallel_apply(replicas, inputs, kwargs)
return self.gather(outputs, self.output_device)
The self.scatter
in the forward function of dataparallel
is trying to port all arguments into the gpu but i dont want inputs2
to be in gpu
Even if you do something to force inputs2
to cpu
then you cannot do operations with that inputs2
because it becomes torch.FloatTensor
and your model parameters will be torch.cuda.FloatTensor
which are incompatible