summaryrefslogtreecommitdiff
path: root/models/layers.py
blob: f3ccbeb55619c60880a374fbb1f5a0e9de1340d4 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
from typing import Union, Tuple

import torch
import torch.nn as nn


class BasicConv2d(nn.Module):
    def __init__(
            self,
            in_channels: int,
            out_channels: int,
            kernel_size: Union[int, Tuple[int, int]],
            **kwargs
    ):
        super(BasicConv2d, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size,
                              bias=False, **kwargs)

    def forward(self, x):
        return self.conv(x)


class VGGConv2d(nn.Module):
    def __init__(
            self,
            in_channels: int,
            out_channels: int,
            kernel_size: Union[int, Tuple[int, int]] = 3,
            padding: int = 1,
            **kwargs
    ):
        super(VGGConv2d, self).__init__()
        self.conv = BasicConv2d(in_channels, out_channels, kernel_size,
                                padding=padding, **kwargs)

    def forward(self, x):
        return self.conv(x)


class BasicConvTranspose2d(nn.Module):
    def __init__(
            self,
            in_channels: int,
            out_channels: int,
            kernel_size: Union[int, Tuple[int, int]],
            **kwargs
    ):
        super(BasicConvTranspose2d, self).__init__()
        self.trans_conv = nn.ConvTranspose2d(in_channels, out_channels,
                                             kernel_size, bias=False, **kwargs)

    def forward(self, x):
        return self.trans_conv(x)


class DCGANConvTranspose2d(nn.Module):
    def __init__(
            self,
            in_channels: int,
            out_channels: int,
            kernel_size: Union[int, Tuple[int, int]] = 4,
            stride: int = 2,
            padding: int = 1,
            **kwargs
    ):
        super(DCGANConvTranspose2d).__init__()
        self.trans_conv = BasicConvTranspose2d(in_channels, out_channels,
                                               kernel_size, stride=stride,
                                               padding=padding, **kwargs)

    def forward(self, x):
        return self.trans_conv(x)


class FocalConv2d(nn.Module):
    def __init__(
            self,
            in_channels: int,
            out_channels: int,
            kernel_size: Union[int, Tuple[int, int]],
            halving: int,
            **kwargs
    ):
        super(FocalConv2d, self).__init__()
        self.halving = halving
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size,
                              bias=False, **kwargs)

    def forward(self, x):
        h = x.size(2)
        split_size = h // 2 ** self.halving
        z = x.split(split_size, dim=2)
        z = torch.cat([self.conv(_) for _ in z], dim=2)
        return z


class BasicConv1d(nn.Module):
    def __init__(
            self,
            in_channels: int,
            out_channels: int,
            kernel_size: Union[int, Tuple[int]],
            **kwargs
    ):
        super(BasicConv1d, self).__init__()
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size,
                              bias=False, **kwargs)

    def forward(self, x):
        return self.conv(x)