Skip to content

Train model on multi GPUs

RachelChen edited this page Jul 20, 2018 · 7 revisions

Check CUDA running: cuda.set_device() cuda.current_device()

Intuitive understanding of Model/Data Parallel

Model Parallel:

device = torch.device("cuda:0")`` gpu_list = list(range(0, torch.cuda.device_count())) `approximator = torch.nn.DataParallel(APXM_edsr(), device_ids=[0, 1, 2, 3, 4]) print("cuda.current_device=", torch.cuda.current_device())`` print(approximator)``

load_state_dict takes a state_dict, not a string path, so you’d need to do something like the following:

model = MyModel((1000,1000)) model.load_state_dict(torch.load(model_path))


Cuda Semantics Official CUDA page

DataParallel source code dataparallel.py

DataParallel user program examples