From f77c638d6dc0ba4533b94fa704e46cf478479aec Mon Sep 17 00:00:00 2001 From: Tanay Wakhare Date: Tue, 23 Jun 2020 14:42:30 -0400 Subject: [PATCH] Pytorch AttentionNet (#9088) --- rllib/models/tests/test_torch_modules.py | 262 ++++++++++++++++++ rllib/models/tf/layers/__init__.py | 6 +- .../models/tf/layers/multi_head_attention.py | 3 +- .../layers/relative_multi_head_attention.py | 2 +- rllib/models/torch/attention_net.py | 227 +++++++++++++++ rllib/models/torch/misc.py | 6 +- rllib/models/torch/modules/__init__.py | 11 + rllib/models/torch/modules/gru_gate.py | 52 ++++ .../torch/modules/multi_head_attention.py | 63 +++++ .../modules/relative_multi_head_attention.py | 133 +++++++++ rllib/models/torch/modules/skip_connection.py | 37 +++ rllib/utils/torch_ops.py | 2 +- 12 files changed, 798 insertions(+), 6 deletions(-) create mode 100644 rllib/models/tests/test_torch_modules.py create mode 100644 rllib/models/torch/modules/__init__.py create mode 100644 rllib/models/torch/modules/gru_gate.py create mode 100644 rllib/models/torch/modules/multi_head_attention.py create mode 100644 rllib/models/torch/modules/relative_multi_head_attention.py create mode 100644 rllib/models/torch/modules/skip_connection.py diff --git a/rllib/models/tests/test_torch_modules.py b/rllib/models/tests/test_torch_modules.py new file mode 100644 index 000000000..c55579642 --- /dev/null +++ b/rllib/models/tests/test_torch_modules.py @@ -0,0 +1,262 @@ +import gym +import numpy as np +import unittest + +from ray.rllib.models.tf.attention_net import relative_position_embedding, \ + GTrXLNet +from ray.rllib.models.tf.layers import MultiHeadAttention +from ray.rllib.models.torch.attention_net import relative_position_embedding \ + as relative_position_embedding_torch, GTrXLNet as TorchGTrXLNet +from ray.rllib.models.torch.modules.multi_head_attention import \ + MultiHeadAttention as TorchMultiHeadAttention +from ray.rllib.utils.framework import try_import_torch, try_import_tf +from ray.rllib.utils.test_utils import framework_iterator + +torch, nn = try_import_torch() +tf = try_import_tf() + + +class TestModules(unittest.TestCase): + """Tests various torch/modules and tf/layers required for AttentionNet""" + + def train_torch_full_model(self, + model, + inputs, + outputs, + num_epochs=250, + state=None, + seq_lens=None): + """Convenience method that trains a Torch model for num_epochs epochs + and tests whether loss decreased, as expected. + + Args: + model (nn.Module): Torch model to be trained. + inputs (torch.Tensor): Training data + outputs (torch.Tensor): Training labels + num_epochs (int): Number of epochs to train for + state (torch.Tensor): Internal state of module + seq_lens (torch.Tensor): Tensor of sequence lengths + """ + + criterion = torch.nn.MSELoss(reduction="sum") + optimizer = torch.optim.Adam(model.parameters(), lr=3e-4) + + # Check that the layer trains correctly + for t in range(num_epochs): + y_pred = model(inputs, state, seq_lens) + loss = criterion(y_pred[0], torch.squeeze(outputs[0])) + + if t % 10 == 1: + print(t, loss.item()) + + if t == 1: + init_loss = loss.item() + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + final_loss = loss.item() + + # The final loss has decreased, which tests + # that the model is learning from the training data. + self.assertLess(final_loss / init_loss, 0.99) + + def train_torch_layer(self, model, inputs, outputs, num_epochs=250): + """Convenience method that trains a Torch model for num_epochs epochs + and tests whether loss decreased, as expected. + + Args: + model (nn.Module): Torch model to be trained. + inputs (torch.Tensor): Training data + outputs (torch.Tensor): Training labels + num_epochs (int): Number of epochs to train for + """ + criterion = torch.nn.MSELoss(reduction="sum") + optimizer = torch.optim.SGD(model.parameters(), lr=1e-4) + + # Check that the layer trains correctly + for t in range(num_epochs): + y_pred = model(inputs) + loss = criterion(y_pred, outputs) + + if t == 1: + init_loss = loss.item() + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + final_loss = loss.item() + + # The final loss has decreased by a factor of 2, which tests + # that the model is learning from the training data. + self.assertLess(final_loss / init_loss, 0.5) + + def train_tf_model(self, + model, + inputs, + outputs, + num_epochs=250, + minibatch_size=32): + """Convenience method that trains a Tensorflow model for num_epochs + epochs and tests whether loss decreased, as expected. + + Args: + model (tf.Model): Torch model to be trained. + inputs (np.array): Training data + outputs (np.array): Training labels + num_epochs (int): Number of training epochs + batch_size (int): Number of samples in each minibatch + """ + + # Configure a model for mean-squared error loss. + model.compile(optimizer="SGD", loss="mse", metrics=["mae"]) + + hist = model.fit( + inputs, + outputs, + verbose=0, + epochs=num_epochs, + batch_size=minibatch_size).history + init_loss = hist["loss"][0] + final_loss = hist["loss"][-1] + + self.assertLess(final_loss / init_loss, 0.5) + + def test_multi_head_attention(self): + """Tests the MultiHeadAttention mechanism of Vaswani et al.""" + # B is batch size + B = 1 + # D_in is attention dim, L is memory_tau + L, D_in, D_out = 2, 32, 10 + + for fw, sess in framework_iterator( + frameworks=("tfe", "torch", "tf"), session=True): + + # Create a single attention layer with 2 heads + if fw == "torch": + + # Create random Tensors to hold inputs and outputs + x = torch.randn(B, L, D_in) + y = torch.randn(B, L, D_out) + + model = TorchMultiHeadAttention( + in_dim=D_in, out_dim=D_out, num_heads=2, head_dim=32) + + self.train_torch_layer(model, x, y) + + else: # framework is tensorflow or tensorflow-eager + + x = np.random.random((B, L, D_in)) + y = np.random.random((B, L, D_out)) + + inputs = tf.keras.layers.Input(shape=(L, D_in)) + + model = tf.keras.Sequential([ + inputs, + MultiHeadAttention( + out_dim=D_out, num_heads=2, head_dim=32) + ]) + self.train_tf_model(model, x, y) + + def test_attention_net(self): + """Tests the GTrXL. Builds a full AttentionNet and checks + that it trains in a supervised setting.""" + + # Checks that torch and tf embedding matrices are the same + with tf.Session().as_default() as sess: + assert np.allclose( + relative_position_embedding(20, 15).eval(session=sess), + relative_position_embedding_torch(20, 15).numpy()) + + # B is batch size + B = 32 + # D_in is attention dim, L is memory_tau + L, D_in, D_out = 2, 16, 2 + + for fw, sess in framework_iterator( + frameworks=("tfe", "torch", "tf"), session=True): + + # Create a single attention layer with 2 heads + if fw == "torch": + # Create random Tensors to hold inputs and outputs + x = torch.randn(B, L, D_in) + y = torch.randn(B, L, D_out) + + value_labels = torch.randn(B, L, D_in) + memory_labels = torch.randn(B, L, D_out) + + attention_net = TorchGTrXLNet( + observation_space=gym.spaces.Box( + low=float("-inf"), high=float("inf"), shape=(D_in, )), + action_space=gym.spaces.Discrete(D_out), + num_outputs=D_out, + model_config={"max_seq_len": 2}, + name="TestTorchAttentionNet", + num_transformer_units=2, + attn_dim=D_in, + num_heads=2, + memory_tau=L, + head_dim=D_out, + ff_hidden_dim=16, + init_gate_bias=2.0) + + init_state = attention_net.get_initial_state() + + # Get initial state and add a batch dimension. + init_state = [np.expand_dims(s, 0) for s in init_state] + seq_lens_init = torch.full(size=(B, ), fill_value=L) + + # Torch implementation expects a formatted input_dict instead + # of a numpy array as input. + input_dict = {"obs": x} + self.train_torch_full_model( + attention_net, + input_dict, [y, value_labels, memory_labels], + num_epochs=250, + state=init_state, + seq_lens=seq_lens_init) + + else: # Framework is tensorflow or tensorflow-eager. + x = np.random.random((B, L, D_in)) + y = np.random.random((B, L, D_out)) + + value_labels = np.random.random((B, L, 1)) + memory_labels = np.random.random((B, L, D_in)) + + # We need to create (N-1) MLP labels for N transformer units + mlp_labels = np.random.random((B, L, D_in)) + + attention_net = GTrXLNet( + observation_space=gym.spaces.Box( + low=float("-inf"), high=float("inf"), shape=(D_in, )), + action_space=gym.spaces.Discrete(D_out), + num_outputs=D_out, + model_config={"max_seq_len": 2}, + name="TestTFAttentionNet", + num_transformer_units=2, + attn_dim=D_in, + num_heads=2, + memory_tau=L, + head_dim=D_out, + ff_hidden_dim=16, + init_gate_bias=2.0) + model = attention_net.trxl_model + + # Get initial state and add a batch dimension. + init_state = attention_net.get_initial_state() + init_state = [np.tile(s, (B, 1, 1)) for s in init_state] + + self.train_tf_model( + model, [x] + init_state, + [y, value_labels, memory_labels, mlp_labels], + num_epochs=50, + minibatch_size=B) + + +if __name__ == "__main__": + import pytest + import sys + + sys.exit(pytest.main(["-v", __file__])) diff --git a/rllib/models/tf/layers/__init__.py b/rllib/models/tf/layers/__init__.py index 54b82ec1e..5328e1667 100644 --- a/rllib/models/tf/layers/__init__.py +++ b/rllib/models/tf/layers/__init__.py @@ -2,5 +2,9 @@ from ray.rllib.models.tf.layers.gru_gate import GRUGate from ray.rllib.models.tf.layers.relative_multi_head_attention import \ RelativeMultiHeadAttention from ray.rllib.models.tf.layers.skip_connection import SkipConnection +from ray.rllib.models.tf.layers.multi_head_attention import MultiHeadAttention -__all__ = ["GRUGate", "RelativeMultiHeadAttention", "SkipConnection"] +__all__ = [ + "GRUGate", "RelativeMultiHeadAttention", "SkipConnection", + "MultiHeadAttention" +] diff --git a/rllib/models/tf/layers/multi_head_attention.py b/rllib/models/tf/layers/multi_head_attention.py index 074f53c6a..ccc461364 100644 --- a/rllib/models/tf/layers/multi_head_attention.py +++ b/rllib/models/tf/layers/multi_head_attention.py @@ -47,5 +47,6 @@ class MultiHeadAttention(tf.keras.layers.Layer): wmat = tf.nn.softmax(masked_score, axis=2) out = tf.einsum("bijh,bjhd->bihd", wmat, values) - out = tf.reshape(out, tf.concat((tf.shape(out)[:2], [H * D]), axis=0)) + shape = tf.concat([tf.shape(out)[:2], [H * D]], axis=0) + out = tf.reshape(out, shape) return self._linear_layer(out) diff --git a/rllib/models/tf/layers/relative_multi_head_attention.py b/rllib/models/tf/layers/relative_multi_head_attention.py index f9837e5f1..eb9d2f9c9 100644 --- a/rllib/models/tf/layers/relative_multi_head_attention.py +++ b/rllib/models/tf/layers/relative_multi_head_attention.py @@ -113,7 +113,7 @@ class RelativeMultiHeadAttention(tf.keras.layers.Layer): x = tf.pad(x, [[0, 0], [0, 0], [1, 0], [0, 0]]) x = tf.reshape(x, [x_size[0], x_size[2] + 1, x_size[1], x_size[3]]) - x = tf.slice(x, [0, 1, 0, 0], [-1, -1, -1, -1]) + x = x[:, 1:, :, :] x = tf.reshape(x, x_size) return x diff --git a/rllib/models/torch/attention_net.py b/rllib/models/torch/attention_net.py index e69de29bb..dc1453eb6 100644 --- a/rllib/models/torch/attention_net.py +++ b/rllib/models/torch/attention_net.py @@ -0,0 +1,227 @@ +""" +[1] - Attention Is All You Need - Vaswani, Jones, Shazeer, Parmar, + Uszkoreit, Gomez, Kaiser - Google Brain/Research, U Toronto - 2017. + https://arxiv.org/pdf/1706.03762.pdf +[2] - Stabilizing Transformers for Reinforcement Learning - E. Parisotto + et al. - DeepMind - 2019. https://arxiv.org/pdf/1910.06764.pdf +[3] - Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context. + Z. Dai, Z. Yang, et al. - Carnegie Mellon U - 2019. + https://www.aclweb.org/anthology/P19-1285.pdf +""" +import numpy as np + +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.models.torch.modules import GRUGate, \ + RelativeMultiHeadAttention, SkipConnection +from ray.rllib.models.torch.recurrent_net import RecurrentNetwork +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch + +torch, nn = try_import_torch() + + +def relative_position_embedding(seq_length, out_dim): + """Creates a [seq_length x seq_length] matrix for rel. pos encoding. + + Denoted as Phi in [2] and [3]. Phi is the standard sinusoid encoding + matrix. + + Args: + seq_length (int): The max. sequence length (time axis). + out_dim (int): The number of nodes to go into the first Tranformer + layer with. + + Returns: + torch.Tensor: The encoding matrix Phi. + """ + inverse_freq = 1 / (10000**(torch.arange(0, out_dim, 2.0) / out_dim)) + pos_offsets = torch.arange(seq_length - 1, -1, -1) + inputs = pos_offsets[:, None] * inverse_freq[None, :] + return torch.cat((torch.sin(inputs), torch.cos(inputs)), dim=-1) + + +class GTrXLNet(RecurrentNetwork, nn.Module): + """A GTrXL net Model described in [2]. + + This is still in an experimental phase. + Can be used as a drop-in replacement for LSTMs in PPO and IMPALA. + For an example script, see: `ray/rllib/examples/attention_net.py`. + + To use this network as a replacement for an RNN, configure your Trainer + as follows: + + Examples: + >> config["model"]["custom_model"] = GTrXLNet + >> config["model"]["max_seq_len"] = 10 + >> config["model"]["custom_model_config"] = { + >> num_transformer_units=1, + >> attn_dim=32, + >> num_heads=2, + >> memory_tau=50, + >> etc.. + >> } + """ + + def __init__(self, + observation_space, + action_space, + num_outputs, + model_config, + name, + num_transformer_units, + attn_dim, + num_heads, + memory_tau, + head_dim, + ff_hidden_dim, + init_gate_bias=2.0): + """Initializes a GTrXLNet. + + Args: + num_transformer_units (int): The number of Transformer repeats to + use (denoted L in [2]). + attn_dim (int): The input and output dimensions of one Transformer + unit. + num_heads (int): The number of attention heads to use in parallel. + Denoted as `H` in [3]. + memory_tau (int): The number of timesteps to store in each + transformer block's memory M (concat'd over time and fed into + next transformer block as input). + head_dim (int): The dimension of a single(!) head. + Denoted as `d` in [3]. + ff_hidden_dim (int): The dimension of the hidden layer within + the position-wise MLP (after the multi-head attention block + within one Transformer unit). This is the size of the first + of the two layers within the PositionwiseFeedforward. The + second layer always has size=`attn_dim`. + init_gate_bias (float): Initial bias values for the GRU gates (two + GRUs per Transformer unit, one after the MHA, one after the + position-wise MLP). + """ + + super().__init__(observation_space, action_space, num_outputs, + model_config, name) + + nn.Module.__init__(self) + + self.num_transformer_units = num_transformer_units + self.attn_dim = attn_dim + self.num_heads = num_heads + self.memory_tau = memory_tau + self.head_dim = head_dim + self.max_seq_len = model_config["max_seq_len"] + self.obs_dim = observation_space.shape[0] + + # Constant (non-trainable) sinusoid rel pos encoding matrix. + Phi = relative_position_embedding(self.max_seq_len + self.memory_tau, + self.attn_dim) + + self.linear_layer = SlimFC( + in_size=self.obs_dim, out_size=self.attn_dim) + + self.layers = [self.linear_layer] + + # 2) Create L Transformer blocks according to [2]. + for i in range(self.num_transformer_units): + # RelativeMultiHeadAttention part. + MHA_layer = SkipConnection( + RelativeMultiHeadAttention( + in_dim=self.attn_dim, + out_dim=self.attn_dim, + num_heads=num_heads, + head_dim=head_dim, + rel_pos_encoder=Phi, + input_layernorm=True, + output_activation=nn.ReLU), + fan_in_layer=GRUGate(self.attn_dim, init_gate_bias)) + + # Position-wise MultiLayerPerceptron part. + E_layer = SkipConnection( + nn.Sequential( + torch.nn.LayerNorm(self.attn_dim), + SlimFC( + in_size=self.attn_dim, + out_size=ff_hidden_dim, + use_bias=False, + activation_fn=nn.ReLU), + SlimFC( + in_size=ff_hidden_dim, + out_size=self.attn_dim, + use_bias=False, + activation_fn=nn.ReLU)), + fan_in_layer=GRUGate(self.attn_dim, init_gate_bias)) + + # Build a list of all layers in order. + self.layers.extend([MHA_layer, E_layer]) + + # Postprocess GTrXL output with another hidden layer. + self.logits = SlimFC( + in_size=self.attn_dim, + out_size=self.num_outputs, + activation_fn=nn.ReLU) + + # Value function used by all RLlib Torch RL implementations. + self._value_out = None + self.values_out = SlimFC( + in_size=self.attn_dim, out_size=1, activation_fn=None) + + @override(RecurrentNetwork) + def forward_rnn(self, inputs, state, seq_lens): + # To make Attention work with current RLlib's ModelV2 API: + # We assume `state` is the history of L recent observations (all + # concatenated into one tensor) and append the current inputs to the + # end and only keep the most recent (up to `max_seq_len`). This allows + # us to deal with timestep-wise inference and full sequence training + # within the same logic. + state = [torch.from_numpy(item) for item in state] + observations = state[0] + memory = state[1:] + + inputs = torch.reshape(inputs, [1, -1, observations.shape[-1]]) + observations = torch.cat( + (observations, inputs), axis=1)[:, -self.max_seq_len:] + + all_out = observations + for i in range(len(self.layers)): + # MHA layers which need memory passed in. + if i % 2 == 1: + all_out = self.layers[i](all_out, memory=memory[i // 2]) + # Either linear layers or MultiLayerPerceptrons. + else: + all_out = self.layers[i](all_out) + + logits = self.logits(all_out) + self._value_out = self.values_out(all_out) + + memory_outs = all_out[2:] + # If memory_tau > max_seq_len -> overlap w/ previous `memory` input. + if self.memory_tau > self.max_seq_len: + memory_outs = [ + torch.cat( + [memory[i][:, -(self.memory_tau - self.max_seq_len):], m], + axis=1) for i, m in enumerate(memory_outs) + ] + else: + memory_outs = [m[:, -self.memory_tau:] for m in memory_outs] + + T = list(inputs.size())[1] # Length of input segment (time). + + # Postprocessing final output. + logits = logits[:, -T:] + self._value_out = self._value_out[:, -T:] + + return logits, [observations] + memory_outs + + @override(RecurrentNetwork) + def get_initial_state(self): + # State is the T last observations concat'd together into one Tensor. + # Plus all Transformer blocks' E(l) outputs concat'd together (up to + # tau timesteps). + return [np.zeros((self.max_seq_len, self.obs_dim), np.float32)] + \ + [np.zeros((self.memory_tau, self.attn_dim), np.float32) + for _ in range(self.num_transformer_units)] + + @override(ModelV2) + def value_function(self): + return torch.reshape(self._value_out, [-1]) diff --git a/rllib/models/torch/misc.py b/rllib/models/torch/misc.py index 5913ff2b6..27cf4c964 100644 --- a/rllib/models/torch/misc.py +++ b/rllib/models/torch/misc.py @@ -92,13 +92,15 @@ class SlimFC(nn.Module): out_size, initializer=None, activation_fn=None, + use_bias=True, bias_init=0.0): super(SlimFC, self).__init__() layers = [] - linear = nn.Linear(in_size, out_size) + linear = nn.Linear(in_size, out_size, bias=use_bias) if initializer: initializer(linear.weight) - nn.init.constant_(linear.bias, bias_init) + if use_bias is True: + nn.init.constant_(linear.bias, bias_init) layers.append(linear) if activation_fn: layers.append(activation_fn()) diff --git a/rllib/models/torch/modules/__init__.py b/rllib/models/torch/modules/__init__.py new file mode 100644 index 000000000..1b1cb9e8b --- /dev/null +++ b/rllib/models/torch/modules/__init__.py @@ -0,0 +1,11 @@ +from ray.rllib.models.torch.modules.gru_gate import GRUGate +from ray.rllib.models.torch.modules.multi_head_attention import \ + MultiHeadAttention +from ray.rllib.models.torch.modules.relative_multi_head_attention import \ + RelativeMultiHeadAttention +from ray.rllib.models.torch.modules.skip_connection import SkipConnection + +__all__ = [ + "GRUGate", "RelativeMultiHeadAttention", "SkipConnection", + "MultiHeadAttention" +] diff --git a/rllib/models/torch/modules/gru_gate.py b/rllib/models/torch/modules/gru_gate.py new file mode 100644 index 000000000..76d0f9eb3 --- /dev/null +++ b/rllib/models/torch/modules/gru_gate.py @@ -0,0 +1,52 @@ +from ray.rllib.utils.framework import try_import_torch + +torch, nn = try_import_torch() + + +class GRUGate(nn.Module): + """Implements a gated recurrent unit for use in AttentionNet""" + + def __init__(self, dim, init_bias=0., **kwargs): + """ + input_shape (torch.Tensor): dimension of the input + init_bias (int): Bias added to every input to stabilize training + """ + super().__init__(**kwargs) + self._init_bias = init_bias + + # Xavier initialization of torch tensors + self._w_r = torch.zeros(dim, dim) + self._w_z = torch.zeros(dim, dim) + self._w_h = torch.zeros(dim, dim) + + self._u_r = torch.zeros(dim, dim) + self._u_z = torch.zeros(dim, dim) + self._u_h = torch.zeros(dim, dim) + + nn.init.xavier_uniform_(self._w_r) + nn.init.xavier_uniform_(self._w_z) + nn.init.xavier_uniform_(self._w_h) + + nn.init.xavier_uniform_(self._u_r) + nn.init.xavier_uniform_(self._u_z) + nn.init.xavier_uniform_(self._u_h) + + self._bias_z = torch.zeros(dim, ).fill_(self._init_bias) + + def forward(self, inputs, **kwargs): + # Pass in internal state first. + h, X = inputs + + r = torch.tensordot(X, self._w_r, dims=1) + \ + torch.tensordot(h, self._u_r, dims=1) + r = torch.sigmoid(r) + + z = torch.tensordot(X, self._w_z, dims=1) + \ + torch.tensordot(h, self._u_z, dims=1) - self._bias_z + z = torch.sigmoid(z) + + h_next = torch.tensordot(X, self._w_h, dims=1) + \ + torch.tensordot((h * r), self._u_h, dims=1) + h_next = torch.tanh(h_next) + + return (1 - z) * h + z * h_next diff --git a/rllib/models/torch/modules/multi_head_attention.py b/rllib/models/torch/modules/multi_head_attention.py new file mode 100644 index 000000000..ea27ad650 --- /dev/null +++ b/rllib/models/torch/modules/multi_head_attention.py @@ -0,0 +1,63 @@ +""" +[1] - Attention Is All You Need - Vaswani, Jones, Shazeer, Parmar, + Uszkoreit, Gomez, Kaiser - Google Brain/Research, U Toronto - 2017. + https://arxiv.org/pdf/1706.03762.pdf +""" +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.utils.torch_ops import sequence_mask + +torch, nn = try_import_torch() + + +class MultiHeadAttention(nn.Module): + """A multi-head attention layer described in [1].""" + + def __init__(self, in_dim, out_dim, num_heads, head_dim, **kwargs): + """ + in_dim (int): Dimension of input + out_dim (int): Dimension of output + num_heads (int): Number of attention heads + head_dim (int): Output dimension of each attention head + """ + super().__init__(**kwargs) + + # No bias or non-linearity. + self._num_heads = num_heads + self._head_dim = head_dim + self._qkv_layer = SlimFC( + in_size=in_dim, out_size=3 * num_heads * head_dim, use_bias=False) + + self._linear_layer = SlimFC( + in_size=num_heads * head_dim, out_size=out_dim, use_bias=False) + + def forward(self, inputs): + L = list(inputs.size())[1] # length of segment + H = self._num_heads # number of attention heads + D = self._head_dim # attention head dimension + + qkv = self._qkv_layer(inputs) + + queries, keys, values = torch.chunk(input=qkv, chunks=3, dim=-1) + queries = queries[:, -L:] # only query based on the segment + + queries = torch.reshape(queries, [-1, L, H, D]) + keys = torch.reshape(keys, [-1, L, H, D]) + values = torch.reshape(values, [-1, L, H, D]) + + score = torch.einsum("bihd,bjhd->bijh", queries, keys) + score = score / D**0.5 + + # causal mask of the same length as the sequence + mask = sequence_mask(torch.arange(1, L + 1), dtype=score.dtype) + mask = mask[None, :, :, None] + mask = mask.float() + + masked_score = score * mask + 1e30 * (mask - 1.) + wmat = nn.functional.softmax(masked_score, dim=2) + + out = torch.einsum("bijh,bjhd->bihd", wmat, values) + shape = list(out.size())[:2] + [H * D] + # temp = torch.cat(temp2, [H * D], dim=0) + out = torch.reshape(out, shape) + return self._linear_layer(out) diff --git a/rllib/models/torch/modules/relative_multi_head_attention.py b/rllib/models/torch/modules/relative_multi_head_attention.py new file mode 100644 index 000000000..204495517 --- /dev/null +++ b/rllib/models/torch/modules/relative_multi_head_attention.py @@ -0,0 +1,133 @@ +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.utils.torch_ops import sequence_mask + +torch, nn = try_import_torch() + + +class RelativeMultiHeadAttention(nn.Module): + """A RelativeMultiHeadAttention layer as described in [3]. + + Uses segment level recurrence with state reuse. + """ + + def __init__(self, + in_dim, + out_dim, + num_heads, + head_dim, + rel_pos_encoder, + input_layernorm=False, + output_activation=None, + **kwargs): + """Initializes a RelativeMultiHeadAttention nn.Module object. + + Args: + in_dim (int): + out_dim (int): + num_heads (int): The number of attention heads to use. + Denoted `H` in [2]. + head_dim (int): The dimension of a single(!) attention head + Denoted `D` in [2]. + rel_pos_encoder (: + input_layernorm (bool): Whether to prepend a LayerNorm before + everything else. Should be True for building a GTrXL. + output_activation (Optional[tf.nn.activation]): Optional tf.nn + activation function. Should be relu for GTrXL. + **kwargs: + """ + super().__init__(**kwargs) + + # No bias or non-linearity. + self._num_heads = num_heads + self._head_dim = head_dim + + # 3=Query, key, and value inputs. + self._qkv_layer = SlimFC( + in_size=in_dim, out_size=3 * num_heads * head_dim, use_bias=False) + + self._linear_layer = SlimFC( + in_size=num_heads * head_dim, + out_size=out_dim, + use_bias=False, + activation_fn=output_activation) + + self._pos_proj = SlimFC( + in_size=in_dim, out_size=num_heads * head_dim, use_bias=False) + + self._uvar = torch.zeros(num_heads, head_dim) + self._vvar = torch.zeros(num_heads, head_dim) + nn.init.xavier_uniform_(self._uvar) + nn.init.xavier_uniform_(self._vvar) + + self._rel_pos_encoder = rel_pos_encoder + self._input_layernorm = None + + if input_layernorm: + self._input_layernorm = torch.nn.LayerNorm(in_dim) + + def forward(self, inputs, memory=None): + T = list(inputs.size())[1] # length of segment (time) + H = self._num_heads # number of attention heads + d = self._head_dim # attention head dimension + + # Add previous memory chunk (as const, w/o gradient) to input. + # Tau (number of (prev) time slices in each memory chunk). + Tau = list(memory.shape)[1] if memory is not None else 0 + if memory is not None: + memory.requires_grad_(False) + inputs = torch.cat((memory, inputs), dim=1) + + # Apply the Layer-Norm. + if self._input_layernorm is not None: + inputs = self._input_layernorm(inputs) + + qkv = self._qkv_layer(inputs) + + queries, keys, values = torch.chunk(input=qkv, chunks=3, dim=-1) + # Cut out Tau memory timesteps from query. + queries = queries[:, -T:] + + queries = torch.reshape(queries, [-1, T, H, d]) + keys = torch.reshape(keys, [-1, T + Tau, H, d]) + values = torch.reshape(values, [-1, T + Tau, H, d]) + + R = self._pos_proj(self._rel_pos_encoder) + R = torch.reshape(R, [T + Tau, H, d]) + + # b=batch + # i and j=time indices (i=max-timesteps (inputs); j=Tau memory space) + # h=head + # d=head-dim (over which we will reduce-sum) + score = torch.einsum("bihd,bjhd->bijh", queries + self._uvar, keys) + pos_score = torch.einsum("bihd,jhd->bijh", queries + self._vvar, R) + score = score + self.rel_shift(pos_score) + score = score / d**0.5 + + # causal mask of the same length as the sequence + mask = sequence_mask( + torch.arange(Tau + 1, T + Tau + 1), dtype=score.dtype) + mask = mask[None, :, :, None] + + masked_score = score * mask + 1e30 * (mask.to(torch.float32) - 1.) + wmat = nn.functional.softmax(masked_score, dim=2) + + out = torch.einsum("bijh,bjhd->bihd", wmat, values) + shape = list(out.shape)[:2] + [H * d] + out = torch.reshape(out, shape) + + return self._linear_layer(out) + + @staticmethod + def rel_shift(x): + # Transposed version of the shift approach described in [3]. + # https://github.com/kimiyoung/transformer-xl/blob/ + # 44781ed21dbaec88b280f74d9ae2877f52b492a5/tf/model.py#L31 + x_size = list(x.shape) + + x = torch.nn.functional.pad(x, (0, 0, 1, 0, 0, 0, 0, 0)) + x = torch.reshape(x, [x_size[0], x_size[2] + 1, x_size[1], x_size[3]]) + x = x[:, 1:, :, :] + x = torch.reshape(x, x_size) + + return x diff --git a/rllib/models/torch/modules/skip_connection.py b/rllib/models/torch/modules/skip_connection.py new file mode 100644 index 000000000..9c85eea85 --- /dev/null +++ b/rllib/models/torch/modules/skip_connection.py @@ -0,0 +1,37 @@ +from ray.rllib.utils.framework import try_import_torch + +torch, nn = try_import_torch() + + +class SkipConnection(nn.Module): + """Skip connection layer. + + Adds the original input to the output (regular residual layer) OR uses + input as hidden state input to a given fan_in_layer. + """ + + def __init__(self, layer, fan_in_layer=None, add_memory=False, **kwargs): + """Initializes a SkipConnection nn Module object. + + Args: + layer (nn.Module): Any layer processing inputs. + fan_in_layer (Optional[nn.Module]): An optional + layer taking two inputs: The original input and the output + of `layer`. + """ + super().__init__(**kwargs) + self._layer = layer + self._fan_in_layer = fan_in_layer + + def forward(self, inputs, **kwargs): + # del kwargs + outputs = self._layer(inputs, **kwargs) + # Residual case, just add inputs to outputs. + if self._fan_in_layer is None: + outputs = outputs + inputs + # Fan-in e.g. RNN: Call fan-in with `inputs` and `outputs`. + else: + # NOTE: In the GRU case, `inputs` is the state input. + outputs = self._fan_in_layer((inputs, outputs)) + + return outputs diff --git a/rllib/utils/torch_ops.py b/rllib/utils/torch_ops.py index f72ff070a..4ad87e3c8 100644 --- a/rllib/utils/torch_ops.py +++ b/rllib/utils/torch_ops.py @@ -77,7 +77,7 @@ def sequence_mask(lengths, maxlen=None, dtype=None): 39036). """ if maxlen is None: - maxlen = lengths.max() + maxlen = int(lengths.max()) mask = ~(torch.ones((len(lengths), maxlen)).to( lengths.device).cumsum(dim=1).t() > lengths).t()