3

【pytorch】【进行中】基本配置

 2 years ago
source link: https://www.guofei.site/2019/12/07/torch_setting.html
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.

【pytorch】【进行中】基本配置

2019年12月07日

Author: Guofei

文章归类: 2-6-torch ，文章编号: 261

版权声明：本文作者是郭飞。转载随意，但需要标明原文链接，并通知本人
原文链接：https://www.guofei.site/2019/12/07/torch_setting.html

Edit

GPU

import torch

torch.cuda.is_available() # 返回 True/False 表示GPU是否可用
torch.cuda.device_count() # 可用的GPU数量

转为GPU


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

model.to(device) # 模型
mytensor = my_tensor.to(device) # tensor

model = nn.DataParallel(model)

案例

model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
  print("Let's use", torch.cuda.device_count(), "GPUs!")
  # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
  model = nn.DataParallel(model)

model.to(device)
您的支持将鼓励我继续创作！

report

Recommend

About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK