Module openpack_torch.models.keypoint.stgcn

Ref: https://colab.research.google.com/github/machine-perception-robotics-group/MPRGDeepLearningLectureNotebook/blob/master/15_gcn/03_action_recognition_ST_GCN.ipynb#scrollTo=Vk-AMCVb5jqM

"""Ref: https://colab.research.google.com/github/machine-perception-robotics-group/MPRGDeepLearningLectureNotebook/blob/master/15_gcn/03_action_recognition_ST_GCN.ipynb#scrollTo=Vk-AMCVb5jqM
"""
import numpy as np
import torch
import torch.nn as nn


class SpatialGraphConvLayer(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, Ks: int):
        """Implementation of Spacial Graph Convolution Layer.
        Args:
            in_channels (int): _description_
            out_channels (int): _description_
            Ks (int): _description_
        """
        super().__init__()
        self.Ks = Ks
        self.conv = nn.Conv2d(in_channels=in_channels,
                              out_channels=out_channels * Ks,
                              kernel_size=1)

    def forward(self, x: torch.Tensor, A: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x (torch.Tensor): shape=(N, CH, FRAMES, VERTEX)
            A (torch.Tensor): shape=(Ks, VERTEX, VERTEX)
        Returns:
            torch.Tensor: the same shape as input ``x``.
        """
        x = self.conv(x)
        n, kc, t, v = x.size()
        x = x.view(n, self.Ks, kc // self.Ks, t, v)

        # Apply GraphConv and sum up features.
        x = torch.einsum('nkctv,kvw->nctw', (x, A))
        return x.contiguous()


class TemporalConvLayer(nn.Module):
    def __init__(
        self,
        in_channels: int,
        Kt: int,
        stride: int = 1,
        dropout: float = 0.5,
    ) -> None:
        """Implementation of temporal convolution layer.
        Args:
            in_channels (int): _description_
            Kt (int): kernel size for temporal domain.
            stride (int): stride for temporal domain.
            dropout (float): _description_
        """
        super().__init__()
        self.block = nn.Sequential(
            nn.BatchNorm2d(in_channels),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Conv2d(in_channels,
                      in_channels,
                      (Kt, 1),
                      (stride, 1),
                      ((Kt - 1) // 2, 0)),
            nn.BatchNorm2d(in_channels),
            nn.ReLU(),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x (torch.Tensor): shape=(BATCH, CH, FRAMES, VERTEX)
        Returns:
            torch.Tensor: the same shape as input
        """
        x = self.block(x)
        return x


class STConvBlock(nn.Module):
    """Implementation of Spatial-temporal convolutional block with
    learnable edge.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        Ks: int = None,
        Kt: int = None,
        num_vertex: int = None,
        stride: int = 1,
        dropout=0.5,
    ):
        super().__init__()
        # 空間グラフの畳み込み
        self.sgc = SpatialGraphConvLayer(
            in_channels=in_channels,
            out_channels=out_channels,
            Ks=Ks,
        )

        # Learnable weight matrix M エッジに重みを与えます. どのエッジが重要かを学習します.
        self.M = nn.Parameter(torch.ones((Ks, num_vertex, num_vertex)))
        self.tgc = TemporalConvLayer(out_channels, Kt, stride)

    def forward(self, x, A):
        x = self.sgc(x, A * self.M)
        x = self.tgc(x)
        return x

# -----------------------------------------------------------------------------


class STGCN4Seg(nn.Module):
    """Implementation of ST-GCN for segmentation task.

    """

    def __init__(
            self,
            in_channels: int = None,
            num_classes: int = None,
            Ks: int = None,
            Kt: int = None,
            A: np.ndarray = None,
    ):
        super().__init__()
        A = torch.tensor(A, dtype=torch.float32, requires_grad=False)
        self.register_buffer('A', A)
        A_size = A.size()
        num_vertex = A.size(1)

        # Batch Normalization
        self.bn = nn.BatchNorm1d(in_channels * A_size[1])

        # STConvBlocks
        self.stgc1 = STConvBlock(
            in_channels,
            32,
            Ks=Ks,
            Kt=Kt,
            num_vertex=num_vertex)
        self.stgc2 = STConvBlock(32, 32, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc3 = STConvBlock(32, 32, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc4 = STConvBlock(32, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)

        self.stgc5 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc6 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc7 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc8 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)

        self.stgc9 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc10 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc11 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc12 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)

        # Prediction
        self.fc = nn.Conv2d(64, num_classes, kernel_size=(1, num_vertex))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x (torch.Tensor): shape=(BATCH, IN_CH, FRAMES, VERTEX)
        Returns:
            torch.Tensor: the same shape as the input ``x``.
        """
        # Batch Normalization
        N, C, T, V = x.size()  # batch, channel, frame, node
        x = x.permute(0, 3, 1, 2).contiguous().view(N, V * C, T)
        x = self.bn(x)
        x = x.view(N, V, C, T).permute(0, 2, 3, 1).contiguous()

        # STGC_blocks
        x = self.stgc1(x, self.A)
        x = self.stgc2(x, self.A)
        x = self.stgc3(x, self.A)
        x = self.stgc4(x, self.A)

        x = self.stgc5(x, self.A)
        x = self.stgc6(x, self.A)
        x = self.stgc7(x, self.A)
        x = self.stgc8(x, self.A)

        x = self.stgc9(x, self.A)
        x = self.stgc10(x, self.A)
        x = self.stgc11(x, self.A)
        x = self.stgc12(x, self.A)

        # Prediction
        x = self.fc(x)
        return x

Classes

class STConvBlock (in_channels: int, out_channels: int, Ks: int = None, Kt: int = None, num_vertex: int = None, stride: int = 1, dropout=0.5)

Implementation of Spatial-temporal convolutional block with learnable edge.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Expand source code

class STConvBlock(nn.Module):
    """Implementation of Spatial-temporal convolutional block with
    learnable edge.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        Ks: int = None,
        Kt: int = None,
        num_vertex: int = None,
        stride: int = 1,
        dropout=0.5,
    ):
        super().__init__()
        # 空間グラフの畳み込み
        self.sgc = SpatialGraphConvLayer(
            in_channels=in_channels,
            out_channels=out_channels,
            Ks=Ks,
        )

        # Learnable weight matrix M エッジに重みを与えます. どのエッジが重要かを学習します.
        self.M = nn.Parameter(torch.ones((Ks, num_vertex, num_vertex)))
        self.tgc = TemporalConvLayer(out_channels, Kt, stride)

    def forward(self, x, A):
        x = self.sgc(x, A * self.M)
        x = self.tgc(x)
        return x

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x, A) ‑> Callable[..., Any]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Expand source code

def forward(self, x, A):
    x = self.sgc(x, A * self.M)
    x = self.tgc(x)
    return x

class STGCN4Seg (in_channels: int = None, num_classes: int = None, Ks: int = None, Kt: int = None, A: numpy.ndarray = None)

Implementation of ST-GCN for segmentation task.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Expand source code

class STGCN4Seg(nn.Module):
    """Implementation of ST-GCN for segmentation task.

    """

    def __init__(
            self,
            in_channels: int = None,
            num_classes: int = None,
            Ks: int = None,
            Kt: int = None,
            A: np.ndarray = None,
    ):
        super().__init__()
        A = torch.tensor(A, dtype=torch.float32, requires_grad=False)
        self.register_buffer('A', A)
        A_size = A.size()
        num_vertex = A.size(1)

        # Batch Normalization
        self.bn = nn.BatchNorm1d(in_channels * A_size[1])

        # STConvBlocks
        self.stgc1 = STConvBlock(
            in_channels,
            32,
            Ks=Ks,
            Kt=Kt,
            num_vertex=num_vertex)
        self.stgc2 = STConvBlock(32, 32, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc3 = STConvBlock(32, 32, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc4 = STConvBlock(32, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)

        self.stgc5 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc6 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc7 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc8 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)

        self.stgc9 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc10 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc11 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)
        self.stgc12 = STConvBlock(64, 64, Ks=Ks, Kt=Kt, num_vertex=num_vertex)

        # Prediction
        self.fc = nn.Conv2d(64, num_classes, kernel_size=(1, num_vertex))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x (torch.Tensor): shape=(BATCH, IN_CH, FRAMES, VERTEX)
        Returns:
            torch.Tensor: the same shape as the input ``x``.
        """
        # Batch Normalization
        N, C, T, V = x.size()  # batch, channel, frame, node
        x = x.permute(0, 3, 1, 2).contiguous().view(N, V * C, T)
        x = self.bn(x)
        x = x.view(N, V, C, T).permute(0, 2, 3, 1).contiguous()

        # STGC_blocks
        x = self.stgc1(x, self.A)
        x = self.stgc2(x, self.A)
        x = self.stgc3(x, self.A)
        x = self.stgc4(x, self.A)

        x = self.stgc5(x, self.A)
        x = self.stgc6(x, self.A)
        x = self.stgc7(x, self.A)
        x = self.stgc8(x, self.A)

        x = self.stgc9(x, self.A)
        x = self.stgc10(x, self.A)
        x = self.stgc11(x, self.A)
        x = self.stgc12(x, self.A)

        # Prediction
        x = self.fc(x)
        return x

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Args

x : torch.Tensor

shape=(BATCH, IN_CH, FRAMES, VERTEX)

Returns

torch.Tensor

the same shape as the input x.

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Args:
        x (torch.Tensor): shape=(BATCH, IN_CH, FRAMES, VERTEX)
    Returns:
        torch.Tensor: the same shape as the input ``x``.
    """
    # Batch Normalization
    N, C, T, V = x.size()  # batch, channel, frame, node
    x = x.permute(0, 3, 1, 2).contiguous().view(N, V * C, T)
    x = self.bn(x)
    x = x.view(N, V, C, T).permute(0, 2, 3, 1).contiguous()

    # STGC_blocks
    x = self.stgc1(x, self.A)
    x = self.stgc2(x, self.A)
    x = self.stgc3(x, self.A)
    x = self.stgc4(x, self.A)

    x = self.stgc5(x, self.A)
    x = self.stgc6(x, self.A)
    x = self.stgc7(x, self.A)
    x = self.stgc8(x, self.A)

    x = self.stgc9(x, self.A)
    x = self.stgc10(x, self.A)
    x = self.stgc11(x, self.A)
    x = self.stgc12(x, self.A)

    # Prediction
    x = self.fc(x)
    return x

class SpatialGraphConvLayer (in_channels: int, out_channels: int, Ks: int)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Implementation of Spacial Graph Convolution Layer.

Args

in_channels : int

description

out_channels : int

description

Ks : int

description

Expand source code

class SpatialGraphConvLayer(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, Ks: int):
        """Implementation of Spacial Graph Convolution Layer.
        Args:
            in_channels (int): _description_
            out_channels (int): _description_
            Ks (int): _description_
        """
        super().__init__()
        self.Ks = Ks
        self.conv = nn.Conv2d(in_channels=in_channels,
                              out_channels=out_channels * Ks,
                              kernel_size=1)

    def forward(self, x: torch.Tensor, A: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x (torch.Tensor): shape=(N, CH, FRAMES, VERTEX)
            A (torch.Tensor): shape=(Ks, VERTEX, VERTEX)
        Returns:
            torch.Tensor: the same shape as input ``x``.
        """
        x = self.conv(x)
        n, kc, t, v = x.size()
        x = x.view(n, self.Ks, kc // self.Ks, t, v)

        # Apply GraphConv and sum up features.
        x = torch.einsum('nkctv,kvw->nctw', (x, A))
        return x.contiguous()

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x: torch.Tensor, A: torch.Tensor) ‑> torch.Tensor

Args

x : torch.Tensor

shape=(N, CH, FRAMES, VERTEX)

A : torch.Tensor

shape=(Ks, VERTEX, VERTEX)

Returns

torch.Tensor

the same shape as input x.

Expand source code

def forward(self, x: torch.Tensor, A: torch.Tensor) -> torch.Tensor:
    """
    Args:
        x (torch.Tensor): shape=(N, CH, FRAMES, VERTEX)
        A (torch.Tensor): shape=(Ks, VERTEX, VERTEX)
    Returns:
        torch.Tensor: the same shape as input ``x``.
    """
    x = self.conv(x)
    n, kc, t, v = x.size()
    x = x.view(n, self.Ks, kc // self.Ks, t, v)

    # Apply GraphConv and sum up features.
    x = torch.einsum('nkctv,kvw->nctw', (x, A))
    return x.contiguous()

class TemporalConvLayer (in_channels: int, Kt: int, stride: int = 1, dropout: float = 0.5)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Implementation of temporal convolution layer.

Args

in_channels : int

description

Kt : int

kernel size for temporal domain.

stride : int

stride for temporal domain.

dropout : float

description

Expand source code

class TemporalConvLayer(nn.Module):
    def __init__(
        self,
        in_channels: int,
        Kt: int,
        stride: int = 1,
        dropout: float = 0.5,
    ) -> None:
        """Implementation of temporal convolution layer.
        Args:
            in_channels (int): _description_
            Kt (int): kernel size for temporal domain.
            stride (int): stride for temporal domain.
            dropout (float): _description_
        """
        super().__init__()
        self.block = nn.Sequential(
            nn.BatchNorm2d(in_channels),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Conv2d(in_channels,
                      in_channels,
                      (Kt, 1),
                      (stride, 1),
                      ((Kt - 1) // 2, 0)),
            nn.BatchNorm2d(in_channels),
            nn.ReLU(),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x (torch.Tensor): shape=(BATCH, CH, FRAMES, VERTEX)
        Returns:
            torch.Tensor: the same shape as input
        """
        x = self.block(x)
        return x

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Args

x : torch.Tensor

shape=(BATCH, CH, FRAMES, VERTEX)

Returns

torch.Tensor

the same shape as input

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Args:
        x (torch.Tensor): shape=(BATCH, CH, FRAMES, VERTEX)
    Returns:
        torch.Tensor: the same shape as input
    """
    x = self.block(x)
    return x