class TemporalConvNet(nn.Module):
def __init__(self, N, B, H, P, X, R, C, norm_type="gLN", causal=False,
mask_nonlinear='relu'):
super(TemporalConvNet, self).__init__()
# Hyper-parameter
self.C = C
self.X = X
# self.z = []
self.mask_nonlinear = mask_nonlinear
# Components
# [M, N, K] -> [M, N, K]
self.layer_norm = ChannelwiseLayerNorm(N)
# [M, N, K] -> [M, B, K]
self.bottleneck_conv1x1 = nn.Conv1d(N, B, 1, bias=False)
# [M, B, K] -> [M, N, K]
self.dwsConv={}
self.conv1x1={}
self.prelu={}
self.norm={}
self.pwConv={}
for d in range(X):
self.dwsConv["depthwise_conv{}".format(d)] = nn.Conv1d(N, N, P,
stride=1, padding=2**d,
dilation=2**d, groups=N,
bias=False)
self.conv1x1["conv1d{}".format(d)] = nn.Conv1d(B, N, 1, bias=False)
self.prelu["prelu{}".format(d)] = nn.PReLU()
self.prelu["prelu{}".format(d+100)] = nn.PReLU()
self.norm["norm{}".format(d)] = chose_norm(norm_type, N)
self.norm["norm{}".format(d+100)] = chose_norm(norm_type, N)
# [M, N, K] -> [M, B, K]
self.pwConv["pointwise_conv{}".format(d)] = nn.Conv1d(N, B, 1, bias=False)
# [M, B, K] -> [M, C*N, K]
self.mask_conv1x1 = nn.Conv1d(B, C*N, 1, bias=False)
def forward(self, x):
"""
Keep this API same with TasNet
Args:
mixture_w: [M, N, K], M is batch size
returns:
est_mask: [M, C, N, K]
"""
M, N, K = x.size()
x = self.layer_norm(x)
x = self.bottleneck_conv1x1(x)
for d in range(self.X):
y = self.conv1x1["conv1d{}".format(d)](x)
y = self.prelu["prelu{}".format(d)](y)
y = self.norm["norm{}".format(d)](y)
y = self.dwsConv["depthwise_conv{}".format(d)]
y = self.prelu["prelu{}".format(d+100)](y)
y = self.norm["norm{}".format(d+100)](y)
y = self.pwConv["pointwise_conv{}".format(d)](y)
x = x + y
score = self.mask_conv1x1(x)
# score = self.network(mixture_w) # [M, N, K] -> [M, C*N, K]
score = score.view(M, self.C, N, K) # [M, C*N, K] -> [M, C, N, K]
if self.mask_nonlinear == 'softmax':
est_mask = F.softmax(score, dim=1)
elif self.mask_nonlinear == 'relu':
est_mask = F.relu(score)
else:
raise ValueError("Unsupported mask non-linear function")
return est_mask
In the code above I removed sequential layers and replaced them with dictionary as it was difficult to check output at each layer.
After removing nn.sequential, newly initialized layer weights are not considered as same sub-module, I guess. They are initialized as non-CUDA tensors.
Whats the solution? how to initialize them as torch.cuda.FloatTensor?
Is there any other way to remove sequential for the code above?
Thanks in Advance.