最近,我在做MAE相关代码实验的时候,网络的构架中有用到nn.Embedding()这个函数,刚开始,我用cpu跑训练的时候是没有报错的,但是cpu跑太慢了,我就改成gpu,结果出现了如下错误:
'''改->'''
device = torch.device("cpu")->device = torch.device("cuda")
'''结果'''
mask_tokens = mask_tokens + self.decoder_pos_emb(mask_index)
...
return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: Expected object of device type cuda but got
device type cpu for argument #3 'index' in call to _th_index_select
根据错误提示,我定位到了mask_tokens = mask_tokens+self.decoder_pos_emb(mask_index)这段代码,并且是self.decoder_pos_emb(mask_index)这个函数错误,因为只有self.decoder_pos_emb()这个函数涉及到Embedding():
self.decoder_pos_emb(mask_index)
self.decoder_pos_emb = nn.Embedding(self.encoder.dec_total_num, decoder_dim)
#构造一个可学习的参数,如果括号中的参数为(64,512),那么就构造一个64行,512列的tensor矩阵
mask_index = mask_index.type(torch.LongTensor)
#mask_index是需要对Embedding矩阵进行掩码的矩阵序列,并将其转化成LongTensor-int64位
#(Embedding函数要求)
我觉得self.decoder_pos_emb(mask_index)应该是没错的,所以应该是传入的参数mask_index出错,所以我出查了查torch.LongTensor参数问题。查了半天,终于找到了问题所在,原来cpu->torch.LongTensor, gpu->torch.cuda.LongTensor,具体可查找官网:
所以只要把数据类型torch.LongTensor改成torch.cuda.LongTensor就没有问题了:
mask_index = mask_index.type(torch.cuda.LongTensor)