spleeter-pytorch/spleeter_pytorch/unet.py at 475be5d6f8238190df40053d0ef67860e788a3f7 · fwcd/spleeter-pytorch · GitHub

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
import torch
import torch.nn.functional as F

from torch import nn

class CustomPad(nn.Module):
    def __init__(self, padding_setting=(1, 2, 1, 2)):
        super().__init__()
        self.padding_setting = padding_setting

    def forward(self, x):
        return F.pad(x, self.padding_setting, "constant", 0)


class CustomTransposedPad(nn.Module):
    def __init__(self, padding_setting=(1, 2, 1, 2)):
        super().__init__()
        self.padding_setting = padding_setting

    def forward(self, x):
        l,r,t,b = self.padding_setting
        return x[:,:,l:-r,t:-b]


def down_block(in_filters, out_filters):
    return nn.Sequential(CustomPad(),
                         nn.Conv2d(in_filters, out_filters, kernel_size=5, stride=2,padding=0)), \
           nn.Sequential(
               nn.BatchNorm2d(out_filters, track_running_stats=True, eps=1e-3, momentum=0.01),
               nn.LeakyReLU(0.2))


def up_block(in_filters, out_filters, dropout=False):
    layers = [
        nn.ConvTranspose2d(in_filters, out_filters, kernel_size=5,stride=2),
        CustomTransposedPad(),
        nn.ReLU(),
        nn.BatchNorm2d(out_filters, track_running_stats=True, eps=1e-3, momentum=0.01)
    ]
    if dropout:
        layers.append(nn.Dropout(0.5))

    return nn.Sequential(*layers)


class UNet(nn.Module):
    def __init__(self, in_channels=2):
        super().__init__()
        self.down1_conv, self.down1_act = down_block(in_channels, 16)
        self.down2_conv, self.down2_act = down_block(16, 32)
        self.down3_conv, self.down3_act = down_block(32, 64)
        self.down4_conv, self.down4_act = down_block(64, 128)
        self.down5_conv, self.down5_act = down_block(128, 256)
        self.down6_conv, self.down6_act = down_block(256, 512)

        self.up1 = up_block(512, 256, dropout=True)
        self.up2 = up_block(512, 128, dropout=True)
        self.up3 = up_block(256, 64, dropout=True)
        self.up4 = up_block(128, 32)
        self.up5 = up_block(64, 16)
        self.up6 = up_block(32, 1)
        self.up7 = nn.Sequential(
            nn.Conv2d(1, in_channels, kernel_size=4, dilation=2, padding=3),
            nn.Sigmoid()
        )

    def forward(self, x):
        d1_conv = self.down1_conv(x)
        d1 = self.down1_act(d1_conv)

        d2_conv = self.down2_conv(d1)
        d2 = self.down2_act(d2_conv)

        d3_conv = self.down3_conv(d2)
        d3 = self.down3_act(d3_conv)

        d4_conv = self.down4_conv(d3)
        d4 = self.down4_act(d4_conv)

        d5_conv = self.down5_conv(d4)
        d5 = self.down5_act(d5_conv)

        d6_conv = self.down6_conv(d5)
        d6 = self.down6_act(d6_conv)

        u1 = self.up1(d6_conv)
        u2 = self.up2(torch.cat([d5_conv, u1], axis=1))
        u3 = self.up3(torch.cat([d4_conv, u2], axis=1))
        u4 = self.up4(torch.cat([d3_conv, u3], axis=1))
        u5 = self.up5(torch.cat([d2_conv, u4], axis=1))
        u6 = self.up6(torch.cat([d1_conv, u5], axis=1))
        u7 = self.up7(u6)
        return u7 * x


if __name__ == '__main__':
    net = UNet(14)
    print(net(torch.rand(1, 14, 20, 48)).shape)