-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathFFTNet_dilconv.py
135 lines (105 loc) · 5.39 KB
/
FFTNet_dilconv.py
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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon May 7 09:46:10 2018
@author: sungkyun
FFTNet model using 2x1 dil-conv
"""
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
# Models with Preset (for convenience)
'''
dim_input: dimension of input (256 for 8-bit mu-law input)
num_layer: number of layers (11 in paper). receptive field = 2^11 (2,048)
io_ch: number of input(=output) channels in each fft layers
skip_ch: number of skip-channels, only required for fft-residual net.
Annotations:
B: batch dimension
C: channel dimension
L: length dimension
'''
def fftnet_base(input_dim=256, num_layer=11, io_ch=256):
return FFTNet(input_dim=input_dim, num_layer=num_layer, io_ch=io_ch, skip_ch=0, bias=True)
def fftnet_residual(input_dim=256, num_layer=11, io_ch=256, skip_ch=256):
return FFTNet(input_dim=input_dim, num_layer=num_layer, io_ch=io_ch, skip_ch=skip_ch, bais=True)
# FFT_Block: define a basic FFT Block
'''
FFT_Block:
- using 2x1 dilated-conv, instead of LR split 1x1 conv.
- described in the paper, section 2.2.
- in case of the first layer used in the first FFT_Block,
we use nn.embedding layer for one-hot index(0-255) entries.
'''
class FFT_Block(nn.Module):
def __init__(self, cond_dim=26, io_ch=int, recep_sz=int, bias=True):
super(FFT_Block, self).__init__()
self.cond_dim=cond_dim # Number of dimensions of condition input
self.io_ch = io_ch
self.recep_sz = recep_sz # Size of receptive field: i.e., the 1st layer has receptive field of 2^11(=2,048). 2nd has 2^10.
self.bias = bias # If True, use bias in 1x1 conv.
self.dilation = int(recep_sz / 2)
self.conv_2x1_LR = nn.Conv1d(in_channels=self.io_ch, out_channels=self.io_ch,
kernel_size=2, stride=1, dilation=self.dilation, bias=self.bias)
self.conv_2x1_VLR = nn.Conv1d(in_channels=self.cond_dim, out_channels=self.io_ch,
kernel_size=2, stride=1, dilation=self.dilation, bias=self.bias)
self.conv_1x1_last = nn.Conv1d(in_channels=self.io_ch, out_channels=self.io_ch,
kernel_size=1, stride=1, bias=self.bias)
return None
def forward(self, x, cond):
z = self.conv_2x1_LR(x) # Eq(1), z = w_L*x_L + w_R*x_R
z = z + self.conv_2x1_VLR(cond) # Eq(2), z = (WL ∗ xL + WR ∗ xR) + (VL ∗ hL + VR ∗ hR)
x = F.relu(self.conv_1x1_last(F.relu(z))) # x = ReLU(conv1x1(ReLU(z)))
return x
'''
FFTNet:
- [11 FFT_blocks] --> [FC_layer] --> [softmax]
'''
class FFTNet(nn.Module):
def __init__(self, input_dim=256, cond_dim=26, num_layer=11, io_ch=256, skip_ch=0, bias=True):
super(FFTNet, self).__init__()
self.input_dim = input_dim # 256 (=num_classes)
self.cond_dim = cond_dim # 26
self.num_layer = num_layer # 11
self.io_ch = io_ch # 256 ch. in the paper
self.skip_ch = skip_ch # Not implemented yet (no skip channel in the paper)
self.bias = bias # If True, use bias in 2x1 conv.
self.max_recep_sz = int(pow(2, self.num_layer)) # 2^11, max receptive field size
# Embedding layer: one-hot_index -> embedding -> 256ch output
self.input_embedding_layer = nn.Embedding(num_embeddings=self.input_dim,
embedding_dim=self.io_ch)
# Constructing FFT Blocks:
blocks = nn.ModuleList()
for l in range(self.num_layer):
recep_sz = int(pow(2, self.num_layer-l)) # 1024, 512, ... 2
blocks.append( FFT_Block(cond_dim=self.cond_dim,
io_ch=self.io_ch,
recep_sz=recep_sz,
bias=self.bias) )
self.fft_blocks=blocks
# Final FC layer:
self.fc = nn.Linear(in_features=self.io_ch, out_features=self.io_ch)
return None
def forward(self, x, cond, gen_mod=False):
# Padding x:
zpad_sz = int(self.max_recep_sz)
x = F.pad(x, (zpad_sz, 0), 'constant', 128) # 128? or 0?
# Embedding(x):
x = self.input_embedding_layer(x) # In : BxL, Out: BxLxC
x = x.permute(0,2,1) # Out: BxCxL
# FFT_Blocks:
for l in range(self.num_layer):
# Padding cond:
zpad_sz = int(self.max_recep_sz/pow(2, l))
padded_cond = F.pad(cond, (zpad_sz, 0), 'constant', 0)
x = self.fft_blocks[l](x, padded_cond)
if gen_mod is True:
x = x[:,:,-1] # In generator mode, take the last one sample only.
x = x.reshape(-1, 1, self.io_ch) # (BxC) --> (Bx1xC)
else:
x = x[:,:,:-1] # In training mode, right-omit 1 is required.
x = x.permute(0,2,1) # (BxCxL) --> (BxLxC)
x = self.fc(x) # (BxLxC)
# NOTE: in PyTorch, softmax() is included in CE loss.
return x