Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Multi gpu training #69

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Multi gpu training #69

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0be67a3
added multi_gpu option to training
psteinb Feb 3, 2020
274dfb9
add sentinals for backend later
psteinb Feb 3, 2020
99e4e55
multi gpu training works for 2D training
psteinb Feb 12, 2020
1b0a8d3
notebook to illustrate use of multiple GPUs
psteinb Feb 12, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
491 changes: 491 additions & 0 deletions examples/2D/denoising2D_BSD68/BSD68_reproducibility_multi_gpu.ipynb
View file
Open in desktop

Large diffs are not rendered by default.

4 changes: 4 additions & 0 deletions n2v/models/n2v_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,8 @@ class N2VConfig(argparse.Namespace):
Learning rate for training. Default: ``0.0004``
train_batch_size : int
Batch size for training. Default: ``16``
train_num_gpus : int
how many gpus to use. Default: ``1``
train_tensorboard : bool
Enable TensorBoard for monitoring training progress. Default: ``True``
train_checkpoint : str
Expand Down Expand Up @@ -137,6 +139,7 @@ def __init__(self, X, **kwargs):
self.train_steps_per_epoch = 400
self.train_learning_rate = 0.0004
self.train_batch_size = 16
self.train_num_gpus = 1
self.train_tensorboard = True
self.train_checkpoint = 'weights_best.h5'
self.train_reduce_lr = {'factor': 0.5, 'patience': 10}
Expand Down Expand Up @@ -211,6 +214,7 @@ def _is_int(v,low=None,high=None):
ok['train_steps_per_epoch'] = _is_int(self.train_steps_per_epoch,1)
ok['train_learning_rate'] = np.isscalar(self.train_learning_rate) and self.train_learning_rate > 0
ok['train_batch_size'] = _is_int(self.train_batch_size,1)
ok['train_num_gpus'] = _is_int(self.train_num_gpus,1)
ok['train_tensorboard'] = isinstance(self.train_tensorboard,bool)
ok['train_checkpoint'] = self.train_checkpoint is None or isinstance(self.train_checkpoint,string_types)
ok['train_reduce_lr'] = self.train_reduce_lr is None or isinstance(self.train_reduce_lr,dict)
Expand Down
14 changes: 14 additions & 0 deletions n2v/models/n2v_standard.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,8 +86,22 @@ def __init__(self, config, name=None, basedir='.'):
self._update_and_check_config()
self._model_prepared = False
self.keras_model = self._build()

if config is None:
self._find_and_load_weights()
else:
# make this a multi_gpu_model
# NB. Does not work with tf 1.14 and 2.2.4 or 2.2.5
if self.config.train_num_gpus > 1:

# https://keras.io/utils/ supported with tf backend only
# TODO: check that K yields tensorflow

from keras.utils import multi_gpu_model
print('Using {} gpus for training'.format(self.config.train_num_gpus))
self.keras_model = multi_gpu_model(self.keras_model,
gpus=self.config.train_num_gpus)



def _build(self):
Expand Down