-
Notifications
You must be signed in to change notification settings - Fork 0
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