From 63496c7bbf41948b1a4b6207bd6cfdc6ff3af48a Mon Sep 17 00:00:00 2001 From: ArneNx Date: Tue, 23 Jun 2020 12:23:16 +0200 Subject: [PATCH] Added move_to_device and options to make cudnn deterministic --- nnfabrik/utility/nn_helpers.py | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/nnfabrik/utility/nn_helpers.py b/nnfabrik/utility/nn_helpers.py index a1f23631..4f837a76 100644 --- a/nnfabrik/utility/nn_helpers.py +++ b/nnfabrik/utility/nn_helpers.py @@ -1,6 +1,9 @@ # helper functions concerning the ANN architecture import torch +from torch import nn +from torch.backends import cudnn + from mlutils.training import eval_state import numpy as np import random @@ -46,8 +49,8 @@ def get_module_output(model, input_shape): :return: output dimensions of the core """ - initial_device = 'cuda' if next(iter(model.parameters())).is_cuda else 'cpu' - device = 'cuda' if torch.cuda.is_available() else 'cpu' + initial_device = "cuda" if next(iter(model.parameters())).is_cuda else "cpu" + device = "cuda" if torch.cuda.is_available() else "cpu" with eval_state(model): with torch.no_grad(): input = torch.zeros(1, *input_shape[1:]).to(device) @@ -64,4 +67,23 @@ def set_random_seed(seed): random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): + cudnn.benchmark = False + cudnn.deterministic = True torch.cuda.manual_seed(seed) + + +def move_to_device(model, gpu=True, multi_gpu=True): + """ + Moves given model to GPU(s) if they are available + :param model: (torch.nn.Module) model to move + :param gpu: (bool) if True attempt to move to GPU + :param multi_gpu: (bool) if True attempt to use multi-GPU + :return: torch.nn.Module, str + """ + device = "cuda" if torch.cuda.is_available() and gpu else "cpu" + if multi_gpu and torch.cuda.device_count() > 1: + print("Using ", torch.cuda.device_count(), "GPUs") + model = nn.DataParallel(model) + model = model.to(device) + return model, device +