pytorch调用gpu进行计算和训练和测试

甜辣uu

于 2021-03-16 15:49:11 发布

阅读量3.3k

点赞数

分类专栏：自然语言处理-nlp-NLP python从入门到精通本文分类

本文链接：https://blog.csdn.net/qq_38735017/article/details/114884700

版权

python从入门到精通同时被 3 个专栏收录

166 篇文章 ¥29.90 ¥99.00

订阅专栏

超级会员免费看

自然语言处理-nlp-NLP

89 篇文章 ¥59.90 ¥99.00

订阅专栏

超级会员免费看

本文分类

10 篇文章

订阅专栏

本文介绍了如何在PyTorch中将训练数据、模型及损失函数迁移到GPU设备上，以利用GPU加速计算。同时强调了在测试阶段同样需要将数据移至GPU进行处理。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

需要把训练数据 x y 和模型损失函数criterion.to(device) 都放到 gpu上模型内部基于tensor 的计算结果不用

用gpu

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

模型定义：
class BertClassificationModel(nn.Module):
    def __init__(self):
        super(BertClassificationModel, self).__init__()   
        model_class, tokenizer_class, pretrained_weights = (tfs.BertModel, tfs.BertTokenizer, 'bert-base-chinese')         
        self.tokenizer = tokenizer_class.from_pretrained(pretrained_weights)
        self.bert = model_class.from_pretrained(pretrained_weights)
        self.dense = nn.Linear(768, 2)  #bert默认的隐藏单元数是768， 输出单元是2，表示二分类
    def forward(self, input_ids,attention_mask):
        bert_output = se