MusicLabeling/model.py at main · Zhang-Zhenning/MusicLabeling · 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
import math
import torch.nn as nn
import torch.nn.parallel
from torch.autograd import Variable


class BasicModel(nn.Module):
    def __init__(self,h,w,in_channel,output_length):
        super(BasicModel,self).__init__()
        self.h = h
        self.w = w
        self.in_channel = in_channel
        self.hidden_channel = in_channel * 2
        self.output_length = output_length

        # we hope both h and w are the power of 2
        assert math.log2(h) == round(math.log2(h)), "spectrogram height must be a power of 2"
        assert math.log2(w) == round(math.log2(w)), "spectrogram width must be a power of 2"

        self.density = int(math.log2(max(h,w)))
        self.kernel_size = (int(4/(self.w/self.h))-1,4-1)
        self.padding_size = (int(1-(self.w != self.h)),1)

        self.backbone = nn.Sequential()
        self.backbone.add_module('input-conv2d',nn.Conv2d(self.in_channel,self.hidden_channel,self.kernel_size,2,self.padding_size,bias=False))
        self.backbone.add_module('input-bn',nn.BatchNorm2d(self.hidden_channel))
        self.backbone.add_module('input-maxPool',nn.MaxPool2d(self.kernel_size,2,self.padding_size))
        self.backbone.add_module('input-relu',nn.ReLU(True))

        for i in range((self.density-3)//2):
            self.backbone.add_module('{0}:inter-{1}-{2}-conv2d'.format(i, self.hidden_channel*(2**i), self.hidden_channel*(2**(i+1))), nn.Conv2d(self.hidden_channel*(2**i),
                                     self.hidden_channel*(2**(i+1)), self.kernel_size, 2, self.padding_size, bias=False))
            self.backbone.add_module(
                '{0}:inter-bn'.format(i), nn.BatchNorm2d(self.hidden_channel*(2**(i+1))))
            self.backbone.add_module(
                '{0}:inter-maxPool'.format(i), nn.MaxPool2d(self.kernel_size, 2, self.padding_size))
            self.backbone.add_module('{0}:inter-relu'.format(i), nn.ReLU(True))

        denses = (self.density-3)//2
        cur_h = (self.h // 4) // (4**denses)
        cur_w = (self.w // 4) // (4**denses)

        self.conv1 = nn.Conv2d(
            self.hidden_channel*(2**((self.density-3)//2)), 2*self.output_length, (cur_h-1, cur_w-1), 2, 0)
        self.classifier = nn.Softmax(dim=1)
        self.loss = nn.MSELoss()
        self.loss_ce = nn.BCELoss()

    def forward(self,input):
        output1 = self.backbone(input)
        output2 = self.conv1(output1)
        output2 = output2.reshape((int(output2.size()[0]),1,self.output_length,2))
        m = torch.squeeze(output2)
        temp = None

        if int(len(m.size())) <= 2:
            m = torch.unsqueeze(m,0)

        for i in range(int(m.size(0))):
            cur = m[i]

            # print(cur.size())
            cur = self.classifier(cur)
            cur = torch.unsqueeze(cur, 0)

            if temp == None:
                temp = cur
                continue
            else:
                temp = torch.cat((temp, cur), dim=0)

        final_output = temp

        return final_output