Convolutional Highway 신경 네트워크

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


class ConvHighWay(nn.Module):
    """
    y = f(x)        ，     
    y = T(x, Wt) * x + (1 - T(x, Wt)) * H(x, Wh)
       highway      ，            。
       ，  x      (B,C,W,H)
        
    https://arxiv.org/abs/1505.00387
    """

    def __init__(self, in_channel, n_layers=1, activation_fn=F.relu):
        super(ConvHighWay, self).__init__()
        self.activation_fn = activation_fn
        self.n_layers = n_layers

        # kernel_size   padding       ，                 
        self.Wh = nn.ModuleList([nn.Conv2d(in_channels=in_channel, out_channels=in_channel, kernel_size=3, padding=1) for _ in range(n_layers)])
        self.Wt = nn.ModuleList([nn.Conv2d(in_channels=in_channel, out_channels=in_channel, kernel_size=3, padding=1) for _ in range(n_layers)])

        #              y = x， bt     ，  sigmoid(Wt * x + bt)   1
        for layer in self.Wt:
            layer.bias.data.fill_(1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for layer_i in range(self.n_layers):
            # H(x, Wh)
            nonlinear_part = self.activation_fn(self.Wh[layer_i](x))
            # T(x, Wt)
            gate = torch.sigmoid(self.Wt[layer_i](x))
            # T(x, Wt) * x + (1 - T(x, Wt)) * H(x, Wh)
            x = gate * x + (1 - gate) * nonlinear_part
        return x


if __name__ == "__main__":
    channel = 3
    highway = ConvHighWay(channel, n_layers=2)
    x = torch.rand((2, channel, 10, 10))

    print(x.size())
    y = highway(x)
    print(y.size())