请问怎么用GPU训练

[qiujunhan]

找了好久资料都不知道怎么用GPU训练，一些设置也改了

-- coding: UTF-8 --

import torch import torch.nn as nn from torch.autograd import Variable import datasets from models import * import torch_util import os, shutil import argparse import test import torchvision import settings

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

Hyper Parameters

num_epochs = 300 batch_size = 20 learning_rate = 0.001

device = torch_util.select_device() device = torch.device("cuda:0")

def main(args): cnn = CNN().to(device)

cnn.train()
criterion = nn.MultiLabelSoftMarginLoss()
optimizer = torch.optim.Adam(cnn.parameters(), lr=learning_rate)

if args.resume:
    cnn.load_state_dict(torch.load(args.model_path, map_location=device))

max_acc = 0
# Train the Model
train_dataloader = datasets.get_train_data_loader()

[qiujunhan]

报错 内容

Traceback (most recent call last): File ".\train.py", line 68, in main(args) File ".\train.py", line 41, in main predict_labels = cnn(images) File "D:\program_temp\anaconda\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl result = self.forward(*input, **kwargs) File "C:\Users\Administrator\Desktop\captcha_identify.pytorch-master\models.py", line 40, in forward out1 = self.layer1(x) File "D:\program_temp\anaconda\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl result = self.forward(*input, **kwargs) File "D:\program_temp\anaconda\lib\site-packages\torch\nn\modules\container.py", line 117, in forward input = module(input) File "D:\program_temp\anaconda\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl result = self.forward(*input, **kwargs) File "D:\program_temp\anaconda\lib\site-packages\torch\nn\modules\conv.py", line 419, in forward return self._conv_forward(input, self.weight) File "D:\program_temp\anaconda\lib\site-packages\torch\nn\modules\conv.py", line 416, in _conv_forward self.padding, self.dilation, self.groups) RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same

[Thresh-Frazier]

train.py中修改：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

cnn = CNN().to(device) images = Variable(images).to(device) labels = Variable(labels.float()).to(device)

![屏幕截图 2021-02-23 103023](https://user-images.githubusercontent.com/21073721/108795942-7b706200-75c2-11eb-8f88-68525acb7768.png