PyTorch 中 GPU 上的数据投射问题

Question

我正在尝试做一个图像分类器，但我在 GPU 上投射数据时遇到问题。

def train(train_loader, net, epoch):

  # Training mode
  net.train()
  
  start = time.time()
  
  epoch_loss  = []
  pred_list, label_list = [], []

  for batch in train_loader:

    #Batch cast on the GPU
    input, label = batch
    input.to(args['device'])
    label.to(args['device'])
    
    #Forward
    ypred = net(input)
    loss = criterion(ypred, label)
    epoch_loss.append(loss.cpu().data)

    _, pred = torch.max(ypred, axis=1) 
    pred_list.append(pred.cpu().numpy())
    label_list.append(label.cpu().numpy())

    #Backward
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
  
  epoch_loss = np.asarray(epoch_loss)
  pred_list  = np.asarray(pred_list).ravel()
  label_list  = np.asarray(label_list).ravel()

  acc = accuracy_score(pred_list, label_list)
  
  end = time.time()
  print('#################### Train ####################')
  print('Epoch %d, Loss: %.4f +/- %.4f, Acc: %.2f, Time: %.2f' % (epoch, epoch_loss.mean(), 
  epoch_loss.std(), acc*100, end-start))
  
  return epoch_loss.mean()


for epoch in range(args['epoch_num']):
  train(train_loader, net, epoch)
  break #Testing

模型已经在 cuda 中，但我收到错误提示

Input type is torch.FloatTensor and not torch.cuda.FloatTensor

input.to(args['device'])有什么问题？

Answer 1

更新： 根据 OP，在火车循环之前的额外 data.to(device) 导致了这个问题。

you are probably getting a string like 0 or cuda from args['device']; you should do this:

'cpu') #pass your args['device'] ``` so then use `device` to move the
model to GPU:  ``` model.to(device) ```

then call the model with:

``` for batch,(data,label) in enumerate(train_loader):

    #Batch cast on the GPU
    data.to(device =device)
    label.to(device =device)

PyTorch 中 GPU 上的数据投射问题

Problem with data cast on the GPU in PyTorch

python

machine-learning

conv-neural-network

pytorch