MCRec_PyTorch/GPUUtil.py at master · Hui-Li/MCRec_PyTorch · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
import torch
import os

device = None
USE_GPU = False

def set_device(gpu_id):
    global device
    global USE_GPU
    if gpu_id == -1:
        device = torch.device("cpu")
    else:
        os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
        if torch.cuda.is_available():
            device = torch.device("cuda")
            USE_GPU = True
        else:
            print("Cannot find GPU id %s! Use CPU." % gpu_id)
            device = torch.device("cpu")

    print("[Setting] device: %s" % device)


def move_to_device(data):
    global device
    global USE_GPU

    if USE_GPU:
        return data.cuda(device, non_blocking=True)
    else:
        return data


def move_model_to_device(model):
    global device
    global USE_GPU

    if USE_GPU:
        model.to(device)