def _jit_compute_loss_fn(predict_fn, loss_fn, n_devices, jit=True): """Returns a (JIT-compiled) function that computes the loss for one step.""" if n_devices == 1: # TODO(lukaszkaiser): remove branch when not needed. def single_compute_loss(opt_state, batch, state, rng): rng, subrng = jax_random.split(rng[0]) loss_val, state = loss_fn(opt_state[0], batch, predict_fn, state, rng) return loss_val, state, [subrng] return backend.jit(single_compute_loss) if jit else single_compute_loss # Else, for n_devices > 1: @functools.partial(backend.pmap, axis_name='batch') def mapped_compute_loss(opt_state, batch, state, rng): """This is a multi-device version of the update function above.""" # We assume all tensors have the first dimension = n_devices. rng, subrng = jax_random.split(rng) loss_val, state = loss_fn(opt_state[0], batch, predict_fn, state, rng) return loss_val, state, subrng def compute_loss(opt_state, batch, state, rng): return mapped_compute_loss(opt_state, _reshape_by_device(batch, n_devices), state, rng) return compute_loss
def __init__(self, model, batch_size, observation_space, action_space, reward_range, discrete_rewards, history_stream, output_dir, model_predict_kwargs=None): """Initializes the env. Args: model: Trax model. batch_size: (int) Number of simulated environments run in parallel. observation_space: (gym.Space) Observation space. action_space: (gym.Space) Action space. reward_range: (tuple) Pair (min_reward, max_reward). discrete_rewards: (bool) Whether to discretize the rewards. history_stream: Iterator yielding batches of initial input data for the model. The format is implementation-specific. output_dir: (str) Output dir. model_predict_kwargs: (dict) Additional model keyword arguments for inference. Useful when different config is needed for training and inference, e.g. train with memory efficient attention and predict with the regular one. """ self._model = model if model_predict_kwargs is None: model_predict_kwargs = {} model_predict = self._model(mode='predict', **model_predict_kwargs) # NOTE: can set non-default PRNG key by: model_predict._set_rng(...) def predict_with_state(*args, **kwargs): output = model_predict(*args, **kwargs) return (output, model_predict.state) self._model_predict = backend.jit(predict_with_state) self._model_initialize = model_predict.initialize_once self._observation_space = observation_space self._action_space = action_space self._reward_range = reward_range self._output_dir = output_dir self._predict_fn = None self._rng = None self._model_state = None self._history_stream = None # Call the super's ctor. It will use some of the member fields, so we call # it in the end. super(SimulatedEnvProblem, self).__init__( batch_size=batch_size, discrete_rewards=discrete_rewards, history_stream=history_stream, ) self.seed()
def _jit_update_fn(predict_fn, loss_fn, optimizer, n_devices, jit=True): """Returns a (JIT-compiled) function that computes updates for one step.""" model_and_loss = tl.Serial(predict_fn, loss_fn) # Gradients are always wrt. the first argument, so putting weights first. def model_and_loss_call(weights, batch, state, rng): res = model_and_loss(batch, weights=weights, state=state, rng=rng) return res, model_and_loss.state if n_devices == 1: # TODO(lukaszkaiser): remove branch when not needed. def single_update(i, opt_state, batch, state, rng): weights, slots, opt_params = opt_state rng, subrng = jax_random.split(rng[0]) grad_fn = backend.grad(model_and_loss_call, has_aux=True) grads, state = grad_fn(weights, batch, state, rng) return optimizer.tree_update(i, grads, weights, slots, opt_params), state, [subrng] return backend.jit(single_update) if jit else single_update # Else, for n_devices > 1: @functools.partial(backend.pmap, axis_name='batch') def mapped_update(i, opt_state, batch, state, rng): """This is a multi-device version of the update function above.""" # We assume all tensors have the first dimension = n_devices. weights, slots, opt_params = opt_state rng, subrng = jax_random.split(rng) grad_fn = backend.grad(model_and_loss_call, has_aux=True) grads, state = grad_fn(weights, batch, state, rng) # We do a psum(1.0) here instead of `n_devices` since `n_devices` is just # the number of devices on this host machine, however psum goes over all # devices of all hosts (ex: a TPU pod) and we need to be averaging over all # of them. grads = jax.tree_util.tree_map( lambda g: backend.psum(g, 'batch') / backend.psum(1.0, 'batch'), grads) return optimizer.tree_update(i, grads, weights, slots, opt_params), state, subrng def update(i, opt_state, batch, state, rng): return mapped_update(np.repeat(i, n_devices), opt_state, batch, state, rng) return update
def _jit_update_fn(predict_fn, loss_fn, optimizer, n_devices, jit=True): """Returns a (JIT-compiled) function that computes updates for one step.""" model_and_loss = layers.Serial(predict_fn, loss_fn) # Gradients are always wrt. the first argument, so putting params first. def model_and_loss_call(params, batch, state, rng): res = model_and_loss(batch, params=params, state=state, rng=rng) return res, model_and_loss.state if n_devices == 1: # TODO(lukaszkaiser): remove branch when not needed. def single_update(i, opt_state, batch, state, rng): params, slots, opt_params = opt_state rng, subrng = jax_random.split(rng[0]) grad_fn = backend.grad(model_and_loss_call, has_aux=True) grads, state = grad_fn(params, batch, state, rng) return optimizer.tree_update(i, grads, params, slots, opt_params), state, [subrng] return backend.jit(single_update) if jit else single_update # Else, for n_devices > 1: @functools.partial(backend.pmap, axis_name='batch') def mapped_update(i, opt_state, batch, state, rng): """This is a multi-device version of the update function above.""" # We assume all tensors have the first dimension = n_devices. params, slots, opt_params = opt_state rng, subrng = jax_random.split(rng) grad_fn = backend.grad(model_and_loss_call, has_aux=True) grads, state = grad_fn(params, batch, state, rng) grads = jax.tree_util.tree_map(lambda g: backend.psum(g, 'batch'), grads) return optimizer.tree_update(i, grads, params, slots, opt_params), state, subrng def update(i, opt_state, batch, state, rng): return mapped_update(np.repeat(i, n_devices), opt_state, batch, state, rng) return update
def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs, output_dir=None, random_seed=None, n_devices=None, save_steps=None, should_save_checkpoints=True, should_write_summaries=True, has_weights=False, nontrainable_param_map=None, mask_id=None, metrics=None): if backend.get_name() == 'jax': self._host_id = jax.host_id() self._host_count = jax.host_count() else: self._host_id = 0 self._host_count = 1 self._is_chief = (self._host_id == 0) if save_steps is None: save_steps = [] self._save_steps = save_steps self._should_save_checkpoints = should_save_checkpoints self._should_write_summaries = should_write_summaries self._has_weights = has_weights self._mask_id = mask_id self._metrics_dict = _METRICS if metrics is None else metrics loss_fn = loss_fn(has_weights=has_weights, mask_id=mask_id) device_count = backend.device_count() n_devices = n_devices or device_count # TODO(lukaszkaiser): remove this restriction when possible. if n_devices != device_count and backend.get_name() == 'jax': raise ValueError( 'JAX cannot work yet with n_devices != all devices: ' '%d != %d' % (n_devices, device_count)) self._n_devices = n_devices # Simple differential seeding of RNG across hosts by host_id and time. if random_seed is None and self._host_count > 1: _, random_seed = divmod( int(time.time() * 1e6) + int(self._host_id * 1e6), 2**32) rng = get_random_number_generator_and_set_seed(random_seed) inputs = inputs(n_devices) self._inputs = inputs # Initialize the learning rate to a dummy value. It will be set in reset(). opt = optimizer(learning_rate=0.0) # Setup the model. model_train = model(mode='train') model_predict_eval = model(mode='eval') # Setup state. rng, init_rng = jax_random.split(rng) self._rngs = np.stack(jax_random.split(rng, n_devices)) first_shape = inputs.input_shape[0] # If the inputs are a tuple/list, add [None] (batch) to each element. if isinstance(first_shape, (list, tuple)): model_input_shape = tuple( tuple([None] + list(shape)) for shape in inputs.input_shape) model_target_shape = tuple( tuple([None] + list(shape)) for shape in inputs.target_shape) else: # Otherwise just add [None] to the input shape. model_input_shape = tuple([None] + list(inputs.input_shape)) model_target_shape = tuple([None] + list(inputs.target_shape)) # Change all None to 1 in input and target shape. model_input_shape = backend.nested_map(lambda x: x or 1, model_input_shape) model_target_shape = backend.nested_map(lambda x: x or 1, model_target_shape) def new_opt_state_and_model_state(input_shape, input_dtype, target_shape, target_dtype, rng): """Returns optimizer and model states suitable for training a model.""" # Combine inputs and targets on the stack. if not isinstance(input_dtype, (list, tuple)): input_dtype = [input_dtype] input_shape = [input_shape] if not isinstance(target_dtype, (list, tuple)): target_dtype = [target_dtype] target_shape = [target_shape] dtypes = list(input_dtype) + list(target_dtype) shapes = list(input_shape) + list(target_shape) if self._has_weights: shapes += list(target_shape) dtypes += [np.float32 for _ in target_dtype] input_signature = tuple( ShapeDtype(s, d) for (s, d) in zip(shapes, dtypes)) # We need to create a new model instance and not reuse `model_train` here, # because `m.initialize` puts cached parameter values in `m` and hence the # next call of `m.initialize` will give wrong results. m = tl.Serial(model(mode='train'), loss_fn) m._set_rng_recursive(rng) # pylint: disable=protected-access weights, state = m.init(input_signature) (slots, opt_params) = opt.tree_init(weights) return (OptState(weights, slots, opt_params), state) if _is_jit_init(): # JIT parameter initialization to avoid memory fragmentation new_opt_state_and_model_state = backend.jit( new_opt_state_and_model_state, static_argnums=(0, 1, 2, 3)) self._new_opt_state_and_model_state = ( lambda: new_opt_state_and_model_state( # pylint: disable=g-long-lambda model_input_shape, self._inputs.input_dtype, model_target_shape, self._inputs.target_dtype, init_rng)) # Arrange and initialize metrics layers. self._metrics = list(sorted(self._metrics_dict.keys())) metrics_layers = [ self._metrics_dict[m](has_weights=self._has_weights, mask_id=self._mask_id) for m in self._metrics ] metrics_in_parallel = tl.Branch(*metrics_layers) # TODO(lukaszkaiser): clean this up once layer API stabilizes. # For now, we need to initialize metric layers somehow, so here we go. # We assume that they do not have any parameters, so this is a dummy. dummy_shapes = ((1, 2), (1, ), (1, )) if self._has_weights else ((1, 2), (1, )) dummy_signature = tuple(ShapeDtype(s) for s in dummy_shapes) metrics_in_parallel._set_rng_recursive(init_rng) # pylint: disable=protected-access m_weights, m_state = metrics_in_parallel.init(dummy_signature) self._metrics_weights = self._for_n_devices(m_weights) self._metrics_state = self._for_n_devices(m_state) # Jit model_predict and update so they're fast. self._jit_eval = _jit_predict_fn(model_predict_eval, metrics_in_parallel, n_devices) self._jit_update_fn = _jit_update_fn(model_train, loss_fn, opt, n_devices) self._model_train = model_train self._model_predict_eval = model_predict_eval self._loss_fn = loss_fn # TODO(pkozakowski): "Learning rate schedules" are currently able to control # control all optimizer parameters and model state, so let's rename them # accordingly. self._lr_schedule = lr_schedule if nontrainable_param_map is None: nontrainable_param_map = {} self._nontrainable_param_map = nontrainable_param_map # Those fields will be set in reset(). self._output_dir = None self._train_sw = None self._eval_sw = None self._history = None self._lr_fn = None self._opt_state = None self._step = None self._model_state = None if output_dir is not None: self.reset(output_dir)