hw_submission(李昀迪): add hw5_20230411

Q1/q1.py

@@ -0,0 +1,193 @@
+"""
+Long Short Term Memory (LSTM) <link https://ieeexplore.ieee.org/abstract/document/6795963 link> is a kind of recurrent neural network that can capture long-short term information.
+This document mainly includes:
+- Pytorch implementation for LSTM.
+- An example to test LSTM.
+For beginners, you can refer to <link https://zhuanlan.zhihu.com/p/32085405 link> to learn the basics about how LSTM works.
+"""
+from typing import Optional, Union, Tuple, List, Dict
+import math
+import torch
+import torch.nn as nn
+from ding.torch_utils import build_normalization
+
+
+class LSTM(nn.Module):
+    """
+    **Overview:**
+        Implementation of LSTM cell with layer norm.
+    """
+
+    def __init__(
+            self,
+            input_size: int,
+            hidden_size: int,
+            num_layers: int,
+            norm_type: Optional[str] = 'LN',
+            dropout: float = 0.
+    ) -> None:
+        # Initialize arguments.
+        super(LSTM, self).__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        # Initialize normalization functions.
+        norm_func = build_normalization(norm_type)
+        self.norm = nn.ModuleList([norm_func(hidden_size * 4) for _ in range(2 * num_layers)])
+        # Initialize LSTM parameters.
+        self.wx = nn.ParameterList()
+        self.wh = nn.ParameterList()
+        dims = [input_size] + [hidden_size] * num_layers
+        for l in range(num_layers):
+            self.wx.append(nn.Parameter(torch.zeros(dims[l], dims[l + 1] * 4)))
+            self.wh.append(nn.Parameter(torch.zeros(hidden_size, hidden_size * 4)))
+        self.bias = nn.Parameter(torch.zeros(num_layers, hidden_size * 4))
+        # Initialize the Dropout Layer.
+        self.use_dropout = dropout > 0.
+        if self.use_dropout:
+            self.dropout = nn.Dropout(dropout)
+        self._init()
+
+    # Dealing with different types of input and return preprocessed prev_state.
+    def _before_forward(self, inputs: torch.Tensor, prev_state: Union[None, List[Dict]]) -> torch.Tensor:
+        seq_len, batch_size = inputs.shape[:2]
+        # If prev_state is None, it indicates that this is the beginning of a sequence. In this case, prev_state will be initialized as zero.
+        if prev_state is None:
+            zeros = torch.zeros(self.num_layers, batch_size, self.hidden_size, dtype=inputs.dtype, device=inputs.device)
+            prev_state = (zeros, zeros)
+        # If prev_state is not None, then preprocess it into one batch.
+        else:
+            assert len(prev_state) == batch_size
+            state = [[v for v in prev.values()] for prev in prev_state]
+            state = list(zip(*state))
+            prev_state = [torch.cat(t, dim=1) for t in state]
+
+        return prev_state
+
+    def _init(self):
+        # Initialize parameters. Each parameter is initialized using a uniform distribution of: $$U(-\sqrt {\frac 1 {HiddenSize}}, -\sqrt {\frac 1 {HiddenSize}})$$
+        gain = math.sqrt(1. / self.hidden_size)
+        for l in range(self.num_layers):
+            torch.nn.init.uniform_(self.wx[l], -gain, gain)
+            torch.nn.init.uniform_(self.wh[l], -gain, gain)
+            if self.bias is not None:
+                torch.nn.init.uniform_(self.bias[l], -gain, gain)
+
+    def forward(
+            self,
+            inputs: torch.Tensor,
+            prev_state: torch.Tensor,
+    ) -> Tuple[torch.Tensor, Union[torch.Tensor, list]]:
+        # The shape of input is: [sequence length, batch size, input size]
+        seq_len, batch_size = inputs.shape[:2]
+        prev_state = self._before_forward(inputs, prev_state)
+
+        H, C = prev_state
+        x = inputs
+        next_state = []
+        for l in range(self.num_layers):
+            h, c = H[l], C[l]
+            new_x = []
+            for s in range(seq_len):
+                # Calculate $$z, z^i, z^f, z^o$$ simultaneously.
+                gate = self.norm[l * 2](torch.matmul(x[s], self.wx[l])
+                                        ) + self.norm[l * 2 + 1](torch.matmul(h, self.wh[l]))
+                if self.bias is not None:
+                    gate += self.bias[l]
+                gate = list(torch.chunk(gate, 4, dim=1))
+                i, f, o, z = gate
+                # $$z^i = \sigma (Wx^ix^t + Wh^ih^{t-1})$$
+                i = torch.sigmoid(i)
+                # $$z^f = \sigma (Wx^fx^t + Wh^fh^{t-1})$$
+                f = torch.sigmoid(f)
+                # $$z^o = \sigma (Wx^ox^t + Wh^oh^{t-1})$$
+                o = torch.sigmoid(o)
+                # $$z = tanh(Wxx^t + Whh^{t-1})$$
+                z = torch.tanh(z)
+                # $$c^t = z^f \odot c^{t-1}+z^i \odot z$$
+                c = f * c + i * z
+                # $$h^t = z^o \odot tanh(c^t)$$
+                h = o * torch.tanh(c)
+                new_x.append(h)
+            next_state.append((h, c))
+            x = torch.stack(new_x, dim=0)
+            # Dropout layer.
+            if self.use_dropout and l != self.num_layers - 1:
+                x = self.dropout(x)
+        next_state = [torch.stack(t, dim=0) for t in zip(*next_state)]
+        # Return list type, split the next_state .
+        h, c = next_state
+        batch_size = h.shape[1]
+        # Split h with shape [num_layers, batch_size, hidden_size] to a list with length batch_size and each element is a tensor with shape [num_layers, 1, hidden_size]. The same operation is performed on c.
+        next_state = [torch.chunk(h, batch_size, dim=1), torch.chunk(c, batch_size, dim=1)]
+        next_state = list(zip(*next_state))
+        next_state = [{k: v for k, v in zip(['h', 'c'], item)} for item in next_state]
+        return x, next_state
+
+
+def pack_data(data: List[torch.Tensor], traj_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Overview:
+        You need to pack variable-length data to regular tensor, return tensor and corresponding mask.
+        If len(data_i) < traj_len, use `null_padding`,
+        else split the whole sequences info different trajectories.
+    Returns:
+        - tensor (:obj:`torch.Tensor`): dtype (torch.float32), shape (traj_len, B, N)
+        - mask (:obj:`torch.Tensor`): dtype (torch.float32), shape (traj_len, B)
+    """
+    new_data = []
+    mask = []
+    for item in data:
+        D, N = item.shape
+        if D < traj_len:
+            null_padding = torch.zeros(traj_len - D, N)
+            new_item = torch.cat([item, null_padding])
+            new_data.append(new_item)
+            item_mask = torch.ones(traj_len)
+            item_mask[D:].zero_()
+            mask.append(item_mask)
+        else:
+            for i in range(0, D, traj_len):
+                item_mask = torch.ones(traj_len)
+                new_item = item[i:i + traj_len]
+                if new_item.shape[0] < traj_len:
+                    new_item = item[-traj_len:]
+                new_data.append(new_item)
+                mask.append(torch.ones(traj_len))
+    new_data = torch.stack(new_data, dim=1)
+    mask = torch.stack(mask, dim=1)
+
+    return new_data, mask
+
+
+def test_lstm():
+    seq_len_list = [32, 49, 24, 78, 45]
+    traj_len = 32
+    N = 10
+    hidden_size = 32
+    num_layers = 2
+
+    variable_len_data = [torch.rand(s, N) for s in seq_len_list]
+    input_, mask = pack_data(variable_len_data, traj_len)
+    assert isinstance(input_, torch.Tensor), type(input_)
+    batch_size = input_.shape[1]
+    assert batch_size == 9, "packed data must have 9 trajectories"
+    lstm = LSTM(N, hidden_size=hidden_size, num_layers=num_layers, norm_type='LN', dropout=0.1)
+
+    prev_state = None
+    for s in range(traj_len):
+        input_step = input_[s:s + 1]
+        output, prev_state = lstm(input_step, prev_state)
+
+    assert output.shape == (1, batch_size, hidden_size)
+    assert len(prev_state) == batch_size
+    assert prev_state[0]['h'].shape == (num_layers, 1, hidden_size)
+    loss = (output * mask.unsqueeze(-1)).mean()
+    loss.backward()
+    for _, m in lstm.named_parameters():
+        assert isinstance(m.grad, torch.Tensor)
+    print('finished')
+
+
+if __name__ == '__main__':
+    test_lstm()