여기에서 제안되지 않았으므로 torch.device올바른 텐서를 초기화 할 때도 매우 편리하기 때문에을 사용하는 방법을 추가하고 device있습니다.

# setting device on GPU if available, else CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)
print()

#Additional Info when using cuda
if device.type == 'cuda':
    print(torch.cuda.get_device_name(0))
    print('Memory Usage:')
    print('Allocated:', round(torch.cuda.memory_allocated(0)/1024**3,1), 'GB')
    print('Cached:   ', round(torch.cuda.memory_cached(0)/1024**3,1), 'GB')

산출:

Using device: cuda

Tesla K80
Memory Usage:
Allocated: 0.3 GB
Cached:    0.6 GB

위에서 언급했듯이 다음을 사용 device하는 것이 가능합니다 .

• 텐서를 각각으로 이동 시키려면 device:

  torch.rand(10).to(device)

• 에 직접 텐서 를 생성 하려면 device:

  torch.rand(10, device=device)

실제 코드를 변경하지 않고도 CPU 와 GPU 간 전환이 편리 합니다.

편집하다:

캐시 및 할당 된 메모리 에 대한 몇 가지 질문과 혼란이 있었 으므로 추가 정보를 추가하고 있습니다.

• torch.cuda.max_memory_cached(device=None)

  주어진 장치에 대해 캐싱 할당자가 관리하는 최대 GPU 메모리를 바이트 단위로 반환합니다.

• torch.cuda.memory_allocated(device=None)

  주어진 장치에 대한 텐서의 현재 GPU 메모리 사용량을 바이트 단위로 반환합니다.

device게시물에서 위에 명시된대로 직접 넘겨 주거나 None을 남겨두면 을 사용합니다 current_device().

답변

$ watch -n 2 nvidia-smi

$ watch -n 3 nvidia-smi --query-gpu=index,gpu_name,memory.total,memory.used,memory.free,temperature.gpu,pstate,utilization.gpu,utilization.memory --format=csv

In [13]: import  torch

In [14]: torch.cuda.is_available()
Out[14]: True

import torch
torch.cuda.is_available()

import torch
dev = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

import torch
dev = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
t1 = torch.randn(1,2)
t2 = torch.randn(1,2).to(dev)
print(t1)  # tensor([[-0.2678,  1.9252]])
print(t2)  # tensor([[ 0.5117, -3.6247]], device='cuda:0')
t1.to(dev)
print(t1)  # tensor([[-0.2678,  1.9252]]) 
print(t1.is_cuda) # False
t1=t1.to(dev)
print(t1)  # tensor([[-0.2678,  1.9252]], device='cuda:0') 
print(t1.is_cuda) # True

class M(nn.Module):
def __init__(self):
    super().__init__()
    self.l1 = nn.Linear(1,2)

def forward(self, x):
    x = self.l1(x)
    return x
model = M()   # not on cuda
model.to(dev) # is on cuda (all parameters)
print(next(model.parameters()).is_cuda) #True

torch.cuda.is_available()

# assuming that 'a' is a tensor created somewhere else
a.device  # returns the device where the tensor is allocated

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

if device.type == 'cuda':
    # do something