Spaces:
Runtime error
Runtime error
File size: 38,519 Bytes
095470e |
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 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 |
from torch import nn
import numpy as np
import torch.nn.functional as F
from torch.nn.utils import weight_norm
import math
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
symbol_length = 73
class GlowTTS(nn.Module):
def __init__(self):
super().__init__()
self.encoder = Encoder()
self.decoder = Decoder()
def forward(self, text, text_len, mel=None, mel_len=None, inference=False):
"""
=====inputs=====
text: (B, T)
text_len: (B) list
mel: (B, 80, F)
mel_len: (B) list
inference: True/False
=====outputs=====
(tuple) (z, z_mean, z_log_std, log_det, z_mask)
z(training) or y(inference): (B, 80, F) | z: latent representation, y: mel-spectrogram
z_mean: (B, 80, F)
z_log_std: (B, 80, F)
log_det: (B) or None
z_mask: (B, 1, F)
(tuple) (x_mean, x_log_std, x_mask)
x_mean: (B, 80, T)
x_log_std: (B, 80, T)
x_mask: (B, 1, T)
(tuple) (attention_alignment, x_log_dur, log_d)
attention_alignment: (B, T, F)
x_log_dur: (B, 1, T) | 추측한 duration의 log scale
log_d: (B, 1, T) | 적절하다고 추측한 alignment에서의 duration의 log scale
"""
x_mean, x_log_std, x_log_dur, x_mask = self.encoder(text, text_len)
# x_std, x_dur 에 log를 붙인 이유는, 논문 저자의 구현에서는 log가 취해진 값으로 간주하기 때문이다.
y, y_len = mel, mel_len
if not inference: # training
y_max_len = y.size(2)
else: # inference
dur = torch.exp(x_log_dur) * x_mask # (B, 1, T)
ceil_dur = torch.ceil(dur) # (B, 1, T)
y_len = torch.clamp_min(torch.sum(ceil_dur, [1, 2]), 1).long() # (B)
# ceil_dur을 [1, 2] 축에 대해 sum한 뒤 최솟값이 1이상이 되도록 설정. 정수 long 타입으로 반환한다.
y_max_len = None
# preprocessing
if y_max_len is not None:
y_max_len = (y_max_len // 2) * 2 # 홀수면 1을 빼서 짝수로 만든다.
y = y[:, :, :y_max_len] # y_max_len에 맞게 y를 조정
y_len = (y_len // 2) * 2 # y_len이 홀수이면 1을 빼서 짝수로 만든다.
# make the z_mask
B = len(y_len)
temp_max = max(y_len)
z_mask = torch.zeros((B, 1, temp_max), dtype=torch.bool).to(device) # (B, 1, F)
for idx, length in enumerate(y_len):
z_mask[idx, :, :length] = True
# make the attention_mask
attention_mask = x_mask.unsqueeze(3) * z_mask.unsqueeze(2) # (B, 1, T, 1) * (B, 1, 1, F) = (B, 1, T, F)
# 주의: Encoder의 attention_mask와는 다른 mask임.
if not inference: # training
z, log_det = self.decoder(y, z_mask, reverse=False)
with torch.no_grad():
x_std_squared_root = torch.exp(-2 * x_log_std) # (B, 80, T)
logp1 = torch.sum(-0.5 * math.log(2 * math.pi) - x_log_std, [1]).unsqueeze(-1) # [(B, T, F)
logp2 = torch.matmul(x_std_squared_root.transpose(1, 2), -0.5 * (z ** 2)) # [(B, T, 80) * (B, 80, F) = (B, T, F)
logp3 = torch.matmul((x_mean * x_std_squared_root).transpose(1,2), z) # (B, T, 80) * (B, 80, F) = (B, T, F)
logp4 = torch.sum(-0.5 * (x_mean ** 2) * x_std_squared_root, [1]).unsqueeze(-1) # (B, T, F)
logp = logp1 + logp2 + logp3 + logp4 # (B, T, F)
"""
logp는 normal distribution N(x_mean, x_std)의 maximum log-likelihood이다.
sum(log(N(z;x_mean, x_std)))를 정규분포 식을 이용하여 분배법칙으로 풀어내면 위와 같은 식이 도출된다.
"""
attention_alignment = maximum_path(logp, attention_mask.squeeze(1)).detach() # alignment (B, T, F)
z_mean = torch.matmul(attention_alignment.transpose(1, 2), x_mean.transpose(1, 2)) # (B, F, T) * (B, T, 80) -> (B, F, 80)
z_mean = z_mean.transpose(1, 2) # (B, 80, F)
z_log_std = torch.matmul(attention_alignment.transpose(1, 2), x_log_std.transpose(1, 2)) # (B, F, T) * (B, T, 80) -> (B, F, 80)
z_log_std = z_log_std.transpose(1, 2) # (B, 80, F)
log_d = torch.log(1e-8 + torch.sum(attention_alignment, -1)).unsqueeze(1) * x_mask # (B, 1, T) | alignment에서 형성된 duration의 log scale
return (z, z_mean, z_log_std, log_det, z_mask), (x_mean, x_log_std, x_mask), (attention_alignment, x_log_dur, log_d)
else: # inference
# generate_path (make attention_alignment using ceil(x_dur))
attention_alignment = generate_path(ceil_dur.squeeze(1), attention_mask.squeeze(1)) # (B, T, F)
z_mean = torch.matmul(attention_alignment.transpose(1, 2), x_mean.transpose(1, 2)) # (B, F, T) * (B, T, 80) -> (B, F, 80)
z_mean = z_mean.transpose(1, 2) # (B, 80, F)
z_log_std = torch.matmul(attention_alignment.transpose(1, 2), x_log_std.transpose(1, 2)) # (B, F, T) * (B, T, 80) -> (B, F, 80)
z_log_std = z_log_std.transpose(1, 2) # (B, 80, F)
log_d = torch.log(1e-8 + torch.sum(attention_alignment, -1)).unsqueeze(1) * x_mask # (B, 1, T) | alignment에서 형성된 duration의 log scale
z = (z_mean + torch.exp(z_log_std) * torch.randn_like(z_mean)) * z_mask # z(latent representation) 생성
y, log_det = self.decoder(z, z_mask, reverse=True) # mel-spectrogram 생성
return (y, z_mean, z_log_std, log_det, z_mask), (x_mean, x_log_std, x_mask), (attention_alignment, x_log_dur, log_d)
##### 아래 논문의 구현이 훨씬 빠르다. 이 논문 구현을 보고 위의 구현을 변경할 필요가 있다. #####
def maximum_path(value, mask, max_neg_val=-np.inf):
""" Numpy-friendly version. It's about 4 times faster than torch version.
value: [b, t_x, t_y]
mask: [b, t_x, t_y]
"""
value = value * mask
device = value.device
dtype = value.dtype
value = value.cpu().detach().numpy()
mask = mask.cpu().detach().numpy().astype(bool)
b, t_x, t_y = value.shape
direction = np.zeros(value.shape, dtype=np.int64)
v = np.zeros((b, t_x), dtype=np.float32)
x_range = np.arange(t_x, dtype=np.float32).reshape(1,-1)
for j in range(t_y):
v0 = np.pad(v, [[0,0],[1,0]], mode="constant", constant_values=max_neg_val)[:, :-1]
v1 = v
max_mask = (v1 >= v0)
v_max = np.where(max_mask, v1, v0)
direction[:, :, j] = max_mask
index_mask = (x_range <= j)
v = np.where(index_mask, v_max + value[:, :, j], max_neg_val)
direction = np.where(mask, direction, 1)
path = np.zeros(value.shape, dtype=np.float32)
index = mask[:, :, 0].sum(1).astype(np.int64) - 1
index_range = np.arange(b)
for j in reversed(range(t_y)):
path[index_range, index, j] = 1
index = index + direction[index_range, index, j] - 1
path = path * mask.astype(np.float32)
path = torch.from_numpy(path).to(device=device, dtype=dtype)
return path
def generate_path(duration, mask):
"""
duration: [b, t_x]
mask: [b, t_x, t_y]
"""
device = duration.device
b, t_x, t_y = mask.shape # (B, T, F)
cum_duration = torch.cumsum(duration, 1) # 누적합, (B, T)
path = torch.zeros(b, t_x, t_y, dtype=mask.dtype).to(device=device) # (B, T, F)
cum_duration_flat = cum_duration.view(b * t_x) # (B*T)
path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype) # (B*T, F)
path = path.view(b, t_x, t_y) # (B, T, F)
path = path.to(torch.float32)
path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:,:-1] # (B, T, F) # T의 차원 맨 앞을 -1한다.
path = path * mask
return path
def sequence_mask(length, max_length=None):
if max_length is None:
max_length = length.max()
x = torch.arange(max_length, dtype=length.dtype, device=length.device)
return x.unsqueeze(0) < length.unsqueeze(1)
def convert_pad_shape(pad_shape):
l = pad_shape[::-1] # [[0, 0], [p, p], [0, 0]]
pad_shape = [item for sublist in l for item in sublist] # [0, 0, p, p, 0, 0]
return pad_shape
def MAS(path, logp, T_max, F_max):
"""
Glow-TTS의 모듈인 maximum_path의 모듈
MAS 알고리즘을 수행하는 함수이다.
=====inputs=====
path: (T, F)
logp: (T, F)
T_max: (1)
F_max: (1)
=====outputs=====
path: (T, F) | 0과 1로 구성된 alignment
"""
neg_inf = -1e9 # negative infinity
# forward
for j in range(F_max):
for i in range(max(0, T_max + j - F_max), min(T_max, j + 1)): # 평행사변형을 생각하라.
# Q_i_j-1 (current)
if i == j:
Q_cur = neg_inf
else:
Q_cur = logp[i, j-1] # j=0이면 i도 0이므로 j-1을 사용해도 된다.
# Q_i-1_j-1 (previous)
if i==0:
if j==0:
Q_prev = 0. # i=0, j=0인 경우에는 logp 값만 반영해야 한다.
else:
Q_prev = neg_inf # i=0인 경우에는 Q_i-1_j-1을 반영하지 않아야 한다.
else:
Q_prev = logp[i-1, j-1]
# logp에 Q를 갱신한다.
logp[i, j] = max(Q_cur, Q_prev) + logp[i, j]
# backtracking
idx = T_max - 1
for j in range(F_max-1, -1, -1): # F_max-1부터 -1까지(-1 포함 없이 0까지) -1씩 감소
path[idx, j] = 1
if idx != 0:
if (logp[idx, j-1] < logp[idx-1, j-1]) or (idx == j):
idx -= 1
return path
def maximum_path(logp, attention_mask):
"""
Glow-TTS에 사용되는 모듈
MAS를 사용하여 alignment를 찾아주는 역할을 한다.
논문 저자 구현에서는 cpython을 이용하여 병렬 처리를 구현한 듯 하나
여기에서는 python만을 이용하여 구현하였다.
=====inputs=====
logp: (B, T, F) | N(x_mean, x_std)의 log-likelihood
attention_mask: (B, T, F)
=====outputs=====
path: (B, T, F) | alignment
"""
B = logp.shape[0]
logp = logp * attention_mask
# 계산은 CPU에서 실행되도록 하기 위해 기존의 device를 저장하고 .cpu().numpy()를 한다.
logp_device = logp.device
logp_type = logp.dtype
logp = logp.data.cpu().numpy().astype(np.float32)
attention_mask = attention_mask.data.cpu().numpy()
path = np.zeros_like(logp).astype(np.int32) # (B, T, F)
T_max = attention_mask.sum(1)[:, 0].astype(np.int32) # (B)
F_max = attention_mask.sum(2)[:, 0].astype(np.int32) # (B)
# MAS 알고리즘
for idx in range(B):
path[idx] = MAS(path[idx], logp[idx], T_max[idx], F_max[idx]) # (T, F)
return torch.from_numpy(path).to(device=logp_device, dtype=logp_type)
def generate_path(ceil_dur, attention_mask):
"""
Glow-TTS에 사용되는 모듈
inference 과정에서 alignment를 만들어낸다.
=====input=====
ceil_dur: (B, T) | 추론한 duration에 ceil 연산한 것 | ex) [[2, 1, 2, 2, ...], [1, 2, 1, 3, ...], ...]
attention_mask: (B, T, F)
=====output=====
path: (B, T, F) | alignment
"""
B, T, Frame = attention_mask.shape
cum_dur = torch.cumsum(ceil_dur, 1)
cum_dur = cum_dur.to(torch.int32) # (B, T) | 누적합 | ex) [[2, 3, 5, 7, ...], [1, 3, 4, 7, ...], ...]
path = torch.zeros(B, T, Frame).to(ceil_dur.device) # (B, T, F) | all False(0)
# make the sequence_mask
for b, batch_cum_dur in enumerate(cum_dur):
for t, each_cum_dur in enumerate(batch_cum_dur):
path[b, t, :each_cum_dur] = torch.ones((1, 1, each_cum_dur)).to(ceil_dur.device)
# cum_dur로부터 True(1)를 path에 새겨넣는다.
path = path - F.pad(path, (0, 0, 1, 0, 0, 0))[:, :-1] # (B, T, F)
"""
ex) batch를 잠시 제외해두고 예시를 든다.
[[1, 1, 0, 0, 0, 0, 0], [[0, 0, 0, 0, 0, 0, 0], [[1, 1, 0, 0, 0, 0, 0],
[1, 1, 1, 0, 0, 0, 0], - [1, 1, 0, 0, 0, 0, 0], = [0, 0, 1, 0, 0, 0, 0],
[1, 1, 1, 1, 1, 0, 0], [1, 1, 1, 0, 0, 0, 0], [0, 0, 0, 1, 1, 0, 0],
[1, 1, 1, 1, 1, 1, 1]] [1, 1, 1, 1, 1, 0, 0]] [0, 0, 0, 0, 0, 1, 1]]
"""
path = path * attention_mask
return path
class Decoder(nn.Module):
def __init__(self):
super().__init__()
self.flows = nn.ModuleList()
for i in range(12):
self.flows.append(ActNorm())
self.flows.append(InvertibleConv())
self.flows.append(AffineCouplingLayer())
def forward(self, x, x_mask, reverse=False):
"""
=====inputs=====
x: (B, 80, F) | mel-spectrogram(Direct) OR latent representation(Reverse)
x_mask: (B, 1, F)
=====outputs=====
z: (B, 80, F) | latent representation(Direct) OR mel-spectrogram(Reverse)
total_log_det: (B) or None | log determinant
"""
if not reverse:
flows = self.flows
total_log_det = 0
else:
flows = reversed(self.flows)
total_log_det = None
x, x_mask = Squeeze(x, x_mask) # (B, 80, F) -> (B, 160, F//2) | (B, 1, F) -> (B, 1, F//2)
for f in flows:
if not reverse:
x, log_det = f(x, x_mask, reverse=reverse)
total_log_det += log_det
else:
x, _ = f(x, x_mask, reverse=reverse)
x, x_mask = Unsqueeze(x, x_mask) # (B, 160, F//2) -> (B, 80, F) | (B, 1, F//2) -> (B, 1, F)
return x, total_log_det
"""
Decoder는 Glow: Generative Flow with Invertible 1×1 Convolutions 논문의 기본 구조를 따라간다.
Glow 논문: https://arxiv.org/pdf/1807.03039.pdf
"""
def Squeeze(x, x_mask):
"""
Decoder의 preprocessing
=====inputs=====
x: (B, 80, F) | mel_spectrogram or latent representation
x_mask: (B, 1, F)
=====outputs=====
x: (B, 160, F//2) | F//2 = [F/2] ([]: 가우스 기호)
x_mask: (B, 160, F//2)
"""
B, C, F = x.size()
x = x[:, :, :(F//2)*2] # F가 홀수이면 맨 뒤 한 frame을 버림.
x = x.view(B, C, F//2, 2) # (B, 80, F//2, 2)
x = x.permute(0, 3, 1, 2).contiguous() # (B, 2, 80, F//2)
x = x.view(B, C*2, F//2) # (B, 160, F//2)
x_mask = x_mask[:, :, 1::2] # (B, 1, F//2) frame을 1부터 한칸씩 건너뛴다.
x = x * x_mask # masking
return x, x_mask
class ActNorm(nn.Module):
"""
Decoder의 1번째 모듈
"""
def __init__(self):
super().__init__()
self.log_s = nn.Parameter(torch.zeros(1, 160, 1)) # Glow 논문의 s에서 log를 취한 것이다. 즉, log[s]
self.bias = nn.Parameter(torch.zeros(1, 160, 1))
def forward(self, x, x_mask, reverse=False):
"""
=====inputs=====
x: (B, 160, F//2) | mel_spectrogram features
x_mask: (B, 1, F//2) | mel_spectrogram features의 mask. (Decoder의 Squeeze에서 변형됨.)
=====outputs=====
z: (B, 160, F//2)
log_det: (B) or None | log_determinant, reverse=True이면 None 반환
"""
x_len = torch.sum(x_mask, [1, 2]) # (B) | 1, 2차원의 값을 더한다. cf. [1, 2] 대신 [2]만 사용하면 shape가 (B, 1)이 된다.
if not reverse:
z = (x * torch.exp(self.log_s) + self.bias) * x_mask # function & masking
log_det = x_len * torch.sum(self.log_s) # log_determinant
# Glow 논문의 Table 1을 확인하라. log_s를 log[s]라 볼 수 있다.
# determinant 대신 log_determinant를 사용하는 이유는 det보다 작은 수치와 적은 계산량 때문으로 추측된다.
else:
z = ((x - self.bias) / torch.exp(self.log_s)) * x_mask # inverse function & masking
log_det = None
return z, log_det
class InvertibleConv(nn.Module):
"""
Decoder의 2번째 모듈
"""
def __init__(self):
super().__init__()
Q = torch.linalg.qr(torch.FloatTensor(4, 4).normal_())[0] # (4, 4)
"""
torch.FloatTensor(4, 4).normal_(): 정규분포 N(0, 1)에서 무작위로 추출한 4x4 matrix
Q, R = torch.linalg.qr(W): QR분해 | Q: 직교 행렬, R: upper traiangular 행렬 cf. det(Q) = 1 or -1
"""
if torch.det(Q) < 0:
Q[:, 0] = -1 * Q[:, 0] # 0번째 열의 부호를 바꿔서 det(Q) = -1로 만든다.
self.W = nn.Parameter(Q)
def forward(self, x, x_mask, reverse=False):
"""
=====inputs=====
x: (B, 160, F//2)
x_mask: (B, 1, F//2)
=====outputs=====
z: (B, 160, F//2)
log_det: (B) or None
"""
B, C, f = x.size() # B, 160, F//2
x_len = torch.sum(x_mask, [1, 2]) # (B)
# channel mixing
x = x.view(B, 2, C//4, 2, f) # (B, 2, 40, 2, F//2)
x = x.permute(0, 1, 3, 2, 4).contiguous() # (B, 2, 2, 40, F//2)
x = x.view(B, 4, C//4, f) # (B, 4, 40, F//2)
# 편의상 log_det부터 구한다.
if not reverse:
weight = self.W
log_det = (C/4) * x_len * torch.logdet(self.W) # (B) | torch.logdet(W): log(det(W))
# height = C/4, width = x_len 인 상황임을 고려하면 Glow 논문의 log_determinant 식과 같다.
else:
weight = torch.linalg.inv(self.W) # inverse matrix
log_det = None
weight = weight.view(4, 4, 1, 1)
z = F.conv2d(x, weight) # (B, 4, 40, F//2) * (4, 4, 1, 1) -> (B, 4, 40, F//2)
"""
F.conv2d(x, weight)의 convolution 연산은 다음과 같이 생각해야 한다.
(B, 4, 40, F//2): (batch_size, in_channels, height, width)
(4, 4, 1, 1): (out_channels, in_channels/groups, kernel_height, kernel_width)
즉, nn.Conv2d(4, 4, kernel_size=(1, 1))인 상황에 가중치를 준 것이다.
"""
# channel unmixing
z = z.view(B, 2, 2, C//4, f) # (B, 4, 40, F//2) -> (B, 2, 2, 40, F//2)
z = z.permute(0, 1, 3, 2, 4).contiguous() # (B, 2, 40, 2, F//2)
z = z.view(B, C, f) * x_mask # (B, 160, F//2) & masking
return z, log_det
class WN(nn.Module):
"""
Decoder의 3번째 모듈인 AffineCouplingLayer의 모듈
해당 구조는 WAVEGLOW: A FLOW-BASED GENERATIVE NETWORK FOR SPEECH SYNTHESIS 로부터 제안되었다.
WaveGlow 논문: https://arxiv.org/pdf/1811.00002.pdf
"""
def __init__(self, dilation_rate=1):
super().__init__()
self.in_layers = nn.ModuleList()
self.res_skip_layers = nn.ModuleList()
for i in range(4):
dilation = dilation_rate ** i # NVIDIA WaveGlow에서는 dilation_rate=2이지만, 여기에서는 1이므로 의미는 없다.
in_layer = weight_norm(nn.Conv1d(192, 2*192, kernel_size=5, dilation=dilation,
padding=((5-1) * dilation)//2)) # (B, 192, F//2) -> (B, 2*192, F//2)
self.in_layers.append(in_layer)
if i < 3:
res_skip_layer = weight_norm(nn.Conv1d(192, 2*192, kernel_size=1)) # (B, 192, F//2) -> (B, 2*192, F//2)
else:
res_skip_layer = weight_norm(nn.Conv1d(192, 192, kernel_size=1)) # (B, 192, F//2) -> (B, 192, F//2)
self.res_skip_layers.append(res_skip_layer)
self.dropout = nn.Dropout(0.05)
def forward(self, x, x_mask):
"""
=====inputs=====
x: (B, 192, F//2)
x_mask: (B, 1, F//2)
=====outputs=====
output: (B, 192, F//2)
"""
output = torch.zeros_like(x) # (B, 192, F//2) all zeros
for i in range(4):
x_in = self.in_layers[i](x) # (B, 192, F//2) -> (B, 2*192, F//2)
x_in = self.dropout(x_in) # dropout
# fused add tanh sigmoid multiply
tanh_act = torch.tanh(x_in[:, :192, :]) # (B, 192, F//2)
sigmoid_act = torch.sigmoid(x_in[:, 192:, :]) # (B, 192, F//2)
acts = sigmoid_act * tanh_act # (B, 192, F//2)
x_out = self.res_skip_layers[i](acts) # (B, 192, F//2) -> (B, 2*192, F//2) or [last](B, 192, F//2)
if i < 3:
x = (x + x_out[:, :192, :]) * x_mask # residual connection & masking
output += x_out[:, 192:, :] # add output
else:
output += x_out # (B, 192, F//2)
output = output * x_mask # masking
return output
class AffineCouplingLayer(nn.Module):
"""
Decoder의 3번째 모듈
"""
def __init__(self):
super().__init__()
self.start_conv = weight_norm(nn.Conv1d(160//2, 192, kernel_size=1)) # (B, 80, F//2) -> (B, 192, F//2)
self.wn = WN()
self.end_conv = nn.Conv1d(192, 160, kernel_size=1) # (B, 192, F//2) -> (B, 160, F//2)
# end_conv의 초기 가중치를 0으로 설정하는 것이 처음에 학습하지 않는 역할을 하며, 이는 학습 안정화에 도움이 된다.
self.end_conv.weight.data.zero_() # weight를 0으로 초기화
self.end_conv.bias.data.zero_() # bias를 0으로 초기화
def forward(self, x, x_mask, reverse=False):
"""
=====inputs=====
x: (B, 160, F//2)
x_mask: (B, 1, F//2)
=====outputs=====
z: (B, 160, F//2)
log_det: (B) or None
"""
B, C, f = x.size() # B, 160, F//2
x_0, x_1 = x[:, :C//2, :], x[:, C//2:, :] # split: (B, 80, F//2) x2
x = self.start_conv(x_0) * x_mask # (B, 80, F//2) -> (B, 192, F//2) & masking
x = self.wn(x, x_mask) # (B, 192, F//2)
out = self.end_conv(x) # (B, 192, F//2) -> (B, 160, F//2)
z_0 = x_0 # (B, 80, F//2)
m = out[:, :C//2, :] # (B, 80, F//2)
log_s = out[:, C//2:, :] # (B, 80, F//2)
if not reverse:
z_1 = (torch.exp(log_s) * x_1 + m) * x_mask # (B, 80, F//2) | function & masking
log_det = torch.sum(log_s * x_mask, [1, 2]) # (B)
else:
z_1 = (x_1 - m) / torch.exp(log_s) * x_mask # (B, 80, F//2) | inverse function & masking
log_det = None
z = torch.cat([z_0, z_1], dim=1) # (B, 160, F//2)
return z, log_det
def Unsqueeze(x, x_mask):
"""
Decoder의 postprocessing
=====inputs=====
x: (B, 160, F//2)
x_mask: (B, 1, F//2)
=====outputs=====
x: (B, 80, F)
x_mask: (B, 1, F)
"""
B, C, f = x.size() # B, 160, F//2
x = x.view(B, 2, C//2, f) # (B, 2, 80, F//2)
x = x.permute(0, 2, 3, 1).contiguous() # (B, 80, F//2, 2)
x = x.view(B, C//2, 2*f) # (B, 160, F)
x_mask = x_mask.unsqueeze(3).repeat(1, 1, 1, 2).view(B, 1, 2*f) # (B, 1, F//2, 1) -> (B, 1, F//2, 2) -> (B, 1, F)
x = x * x_mask # masking
return x, x_mask
class Encoder(nn.Module):
def __init__(self):
super().__init__()
self.embedding = nn.Embedding(symbol_length, 192) # (B, T) -> (B, T, 192)
nn.init.normal_(self.embedding.weight, 0.0, 192**(-0.5)) # 가중치 정규분포 초기화 (N(0, 0.07xx))
self.prenet = PreNet()
self.transformer_encoder = TransformerEncoder()
self.project_mean = nn.Conv1d(192, 80, kernel_size=1) # (B, 192, T) -> (B, 80, T)
self.project_std = nn.Conv1d(192, 80, kernel_size=1) # (B, 192, T) -> (B, 80, T)
self.duration_predictor = DurationPredictor()
def forward(self, text, text_len):
"""
=====inputs=====
text: (B, Max_T)
text_len: (B)
=====outputs=====
x_mean: (B, 80, T) | 평균, 논문 저자 구현의 train.py에서 out_channels를 80으로 설정한 것을 알 수 있음.
x_std: (B, 80, T) | 표준편차
x_dur: (B, 1, T)
x_mask: (B, 1, T)
"""
x = self.embedding(text) * math.sqrt(192) # (B, T) -> (B, T, 192) # math.sqrt(192) = 13.xx (수정)
x = x.transpose(1, 2) # (B, T, 192) -> (B, 192, T)
# Make the x_mask
x_mask = torch.zeros_like(x[:, 0:1, :], dtype=torch.bool) # (B, 1, T)
for idx, length in enumerate(text_len):
x_mask[idx, :, :length] = True
x = self.prenet(x, x_mask) # (B, 192, T)
x = self.transformer_encoder(x, x_mask) # (B, 192, T)
# project
x_mean = self.project_mean(x) * x_mask # (B, 192, T) -> (B, 80, T)
# x_std = self.project_std(x) * x_mask # (B, 192, T) -> (B, 80, T)
##### 아래는 mean_only를 적용한 것임. #####
x_std = torch.zeros_like(x_mean) # x_log_std: (B, 80, T), all zero # log std = 0이므로 std = 1로 계산됨.
# duration predictor
x_dp = torch.detach(x) # stop_gradient
x_dur = self.duration_predictor(x_dp, x_mask) # (B, 192, T) -> (B, 1, T)
return x_mean, x_std, x_dur, x_mask
class LayerNorm(nn.Module):
"""
여러 곳에서 정규화(Norm)를 위해 사용되는 모듈.
nn.LayerNorm이 이미 pytorch 안에 구현되어 있으나, 항상 마지막 차원을 정규화한다.
그래서 channel을 기준으로 정규화하는 LayerNorm을 따로 구현한다.
"""
def __init__(self, channels):
"""
channels: 입력 데이터의 channel 수 | LayerNorm은 channel 차원을 정규화한다.
"""
super().__init__()
self.channels = channels
self.eps = 1e-4
self.gamma = nn.Parameter(torch.ones(channels)) # 학습 가능한 파라미터
self.beta = nn.Parameter(torch.zeros(channels)) # 학습 가능한 파라미터
def forward(self, x):
"""
=====inputs=====
x: (B, channels, *) | 정규화할 입력 데이터
=====outputs=====
x: (B, channels, *) | channel 차원이 정규화된 데이터
"""
mean = torch.mean(x, dim=1, keepdim=True) # channel 차원(index=1)의 평균 계산, 차원을 유지한다.
variance = torch.mean((x-mean)**2, dim=1, keepdim=True) # 분산 계산
x = (x - mean) * (variance + self.eps)**(-0.5) # (x - m) / sqrt(v)
n = len(x.shape)
shape = [1] * n
shape[1] = -1 # shape = [1, -1, 1] or [1, -1, 1, 1]
x = x * self.gamma.view(*shape) + self.beta.view(*shape) # y = x*gamma + beta
return x
class PreNet(nn.Module):
"""
Encoder의 1번째 모듈
"""
def __init__(self):
super().__init__()
self.convs = nn.ModuleList()
self.norms = nn.ModuleList()
self.relu = nn.ReLU()
self.dropout = nn.Dropout(0.5)
for i in range(3):
self.convs.append(nn.Conv1d(192, 192, kernel_size=5, padding=2)) # (B, 192, T) 유지
self.norms.append(LayerNorm(192)) # (B, 192, T) 유지
self.linear = nn.Conv1d(192, 192, kernel_size=1) # (B, 192, T) 유지 | linear 역할을 하는 conv
def forward(self, x, x_mask):
"""
=====inputs=====
x: (B, 192, T) | Embedding된 입력 데이터
x_mask: (B, 1, T) | 글자 길이에 따른 mask (글자가 있으면 True, 없으면 False로 구성)
=====outputs=====
x: (B, 192, T)
"""
x0 = x
for i in range(3):
x = self.convs[i](x * x_mask)
x = self.norms[i](x)
x = self.relu(x)
x = self.dropout(x)
x = self.linear(x)
x = x0 + x # residual connection
return x
class MultiHeadAttention(nn.Module):
"""
Encoder 중 2번째 모듈인 TransformerEncoder의 1번째 모듈
"""
def __init__(self):
super().__init__()
self.n_heads = 2
self.window_size = 4
self.k_channels = 192 // self.n_heads # 96
self.linear_q = nn.Conv1d(192, 192, kernel_size=1) # (B, 192, T) 유지
self.linear_k = nn.Conv1d(192, 192, kernel_size=1) # (B, 192, T) 유지
self.linear_v = nn.Conv1d(192, 192, kernel_size=1) # (B, 192, T) 유지
nn.init.xavier_uniform_(self.linear_q.weight)
nn.init.xavier_uniform_(self.linear_k.weight)
nn.init.xavier_uniform_(self.linear_v.weight)
relative_std = self.k_channels ** (-0.5) # 0.1xx
self.relative_k = nn.Parameter(torch.randn(1, self.window_size * 2 + 1, self.k_channels) * relative_std) # (1, 9, 96)
self.relative_v = nn.Parameter(torch.randn(1, self.window_size * 2 + 1, self.k_channels) * relative_std) # (1, 9, 96)
self.attention_weights = None
self.linear_out = nn.Conv1d(192, 192, kernel_size=1) # (B, 192, T) 유지
self.dropout = nn.Dropout(0.1)
def forward(self, query, context, attention_mask, self_attention=True):
"""
=====inputs=====
query: (B, 192, T_target) | Glow-TTS에서는 self-attention만 이용하므로 query와 context가 동일한 텐서 x이다.
context: (B, 192, T_source) | query = context || 여기에서는 특히 T_source = T_target 이다.
attention_mask: (B, 1, T, T) | x_mask.unsqueeze(2) * z_mask.unsqueeze(3)
self_attention: True/False | self_attention일 때 relative position representations를 적용한다. 여기에서는 항상 True이다.
# 실제로는 query와 context에 같은 텐서 x를 입력하면 된다.
=====outputs=====
output: (B, 192, T)
"""
query = self.linear_q(query)
key = self.linear_k(context)
value = self.linear_v(context)
B, _, T_tar = query.size()
T_src = key.size(2)
query = query.view(B, self.n_heads, self.k_channels, T_tar).transpose(2, 3)
key = key.view(B, self.n_heads, self.k_channels, T_src).transpose(2, 3)
value = value.view(B, self.n_heads, self.k_channels, T_src).transpose(2, 3)
# (B, 192, T_src) -> (B, 2, 96, T_src) -> (B, 2, T_src, 96)
scores = torch.matmul(query, key.transpose(2, 3)) / (self.k_channels ** 0.5)
# (B, 2, T_tar, 96) * (B, 2, 96, T_src) -> (B, 2, T_tar, T_src)
if self_attention: # True
# Get relative embeddings (relative_keys) (1-1)
padding = max(T_src - (self.window_size + 1), 0) # max(T-5, 0)
start_pos = max((self.window_size + 1) - T_src, 0) # max(5-T, 0)
end_pos = start_pos + 2 * T_src - 1 # (2*T-1) or (T+4)
relative_keys = F.pad(self.relative_k, (0, 0, padding, padding))
# (1, 9, 96) -> (1, pad+9+pad, 96) = (1, 2T-1, 96)
"""
위 코드의 F.pad(input, pad) 에서 pad = (0, 0, padding, padding)은 다음을 의미한다.
- 앞의 (0, 0): input의 -1차원을 앞으로 0, 뒤로 0만큼 패딩한다.
- 앞의 (padding, padding): input의 -2차원을 앞으로 padding, 뒤로 padding만큼 패딩한다.
즉, F.pad에서 pad는 역순으로 생각해주어야 한다.
"""
relative_keys = relative_keys[:, start_pos:end_pos, :] # (1, 2T-1, 96)
# Matmul with relative keys (2-1)
relative_keys = relative_keys.unsqueeze(0).transpose(2, 3) # (1, 2T-1, 96) -> (1, 1, 2T-1, 96) -> (1, 1, 96, 2T-1)
x = torch.matmul(query, relative_keys) # (B, 2, T_tar, 96) * (1, 1, 96, 2T_src-1) = (B, 2, T, 2T-1)
# self attention에서는 T_tar = T_src이므로 이를 다르게 고려할 필요가 없다.
# Relative position to absolute position (3-1)
T = T_tar # Absolute position to relative position에서도 쓰임.
x = F.pad(x, (0, 1)) # (B, 2, T, 2*T-1) -> (B, 2, T, 2*T)
x = x.view(B, self.n_heads, T * 2 * T) # (B, 2, T, 2*T) -> (B, 2. 2T^2)
x = F.pad(x, (0, T-1)) # (B, 2, 2T^2 + T - 1)
x = x.view(B, self.n_heads, T+1, 2*T-1) # (B, 2, T+1, 2T-1)
relative_logits = x[:, :, :T, T-1:] # (B, 2, T, T)
# Compute scores
scores_local = relative_logits / (self.k_channels ** 0.5)
scores = scores + scores_local # (B, 2, T, T)
"""
위 식은 Self-Attention with Relative Position Representations 논문의 5번 식을 구현한 것이다.
Relative- 논문: https://arxiv.org/pdf/1803.02155.pdf
"""
scores = scores.masked_fill(attention_mask == 0, -1e-4) # attention_mask가 0인 곳을 -1e-4로 채운다.
attention_weights = F.softmax(scores, dim=-1) # (B, 2, T_tar, T_src) # Relative- 논문에서의 alpha에 해당한다.
attention_weights = self.dropout(attention_weights) # dropout하는 이유가 무엇일까?
output = torch.matmul(attention_weights, value) # (B, 2, T_tar, T_src) * (B, 2, T_src, 96) -> (B, 2, T_tar, 96)
if self_attention: # True
# Absolute position to relative position (3-2)
x = F.pad(attention_weights, (0, T-1)) # (B, 2, T, T) -> (B, 2, T, 2T-1)
x = x.view((B, self.n_heads, T * (2*T-1))) # (B, 2, 2T^2-T)
x = F.pad(x, (T, 0)) # (B, 2, 2T^2) # 앞에 패딩
x = x.view((B, self.n_heads, T, 2*T)) # (B, 2, T, 2T)
relative_weights = x[:, :, :, 1:] # (B, 2, T, 2T-1)
# Get relative embeddings (relative_value) (1-2) # (1-1)과 거의 동일
padding = max(T_src - (self.window_size + 1), 0) # max(T-5, 0)
start_pos = max((self.window_size + 1) - T_src, 0) # max(5-T, 0)
end_pos = start_pos + 2 * T_src - 1 # (2*T-1) or (T+4)
relative_values = F.pad(self.relative_v, (0, 0, padding, padding))
# (1, 9, 96) -> (1, pad+9+pad, 96) = (1, 2T-1, 96)
relative_values = relative_values[:, start_pos:end_pos, :] # (1, 2T-1, 96)
# Matmul with relative values (2-2)
relative_values = relative_values.unsqueeze(0) # (1, 1, 2T-1, 96)
output = output + torch.matmul(relative_weights, relative_values)
# (B, 2, T, 2T-1) * (1, 1, 2T-1, 96) = (B, 2, T, 96)
"""
위 식은 Self-Attention with Relative Position Representations 논문의 3번 식을 구현한 것이다. (분배법칙 이용)
Relative- 논문: https://arxiv.org/pdf/1803.02155.pdf
"""
output = output.transpose(2, 3).contiguous().view(B, 192, T_tar)
# (B, 2, 96, T) -> 메모리에 연속 배치 -> (B, 192, T)
self.attention_weights = attention_weights # (B, 2, T, T)
output = self.linear_out(output)
return output # (B, 192, T)
class FFN(nn.Module):
"""
Encoder 중 2번째 모듈인 TransformerEncoder의 2번째 모듈
"""
def __init__(self):
super().__init__()
self.conv1 = nn.Conv1d(192, 768, kernel_size=3, padding=1) # (B, 192, T) -> (B, 768, T)
self.relu = nn.ReLU()
self.conv2 = nn.Conv1d(768, 192, kernel_size=3, padding=1) # (B, 768, T) -> (B, 192, T)
self.dropout = nn.Dropout(0.1)
def forward(self, x, x_mask):
"""
=====inputs=====
x: (B, 192, T)
x_mask: (B, 1, T)
=====outputs=====
output: (B, 192, T)
"""
x = self.conv1(x)
x = self.relu(x)
x = self.dropout(x)
x = self.conv2(x)
output = x * x_mask
return output
class TransformerEncoder(nn.Module):
"""
Encoder의 2번째 모듈
"""
def __init__(self):
super().__init__()
self.attentions = nn.ModuleList()
self.norms1 = nn.ModuleList()
self.ffns = nn.ModuleList()
self.norms2 = nn.ModuleList()
for i in range(6):
self.attentions.append(MultiHeadAttention())
self.norms1.append(LayerNorm(192))
self.ffns.append(FFN())
self.norms2.append(LayerNorm(192))
self.dropout = nn.Dropout(0.1)
def forward(self, x, x_mask):
"""
=====inputs=====
x: (B, 192, T)
x_mask: (B, 1, T)
=====outputs=====
output: (B, 192, T)
"""
attention_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(3)
# (B, 1, 1, T) * (B, 1, T, 1) = (B, 1, T, T), only consist 0 or 1
for i in range(6):
x = x * x_mask
y = self.attentions[i](x, x, attention_mask)
y = self.dropout(y)
x = x + y # residual connection
x = self.norms1[i](x) # (B, 192, T) 유지
y = self.ffns[i](x, x_mask)
y = self.dropout(y)
x = x + y # residual connection
x = self.norms2[i](x)
output = x * x_mask
return output # (B, 192, T)
class DurationPredictor(nn.Module):
"""
Encoder의 3번째 모듈
"""
def __init__(self):
super().__init__()
self.conv1 = nn.Conv1d(192, 256, kernel_size=3, padding=1) # (B, 192, T) -> (B, 256, T)
self.norm1 = LayerNorm(256)
self.conv2 = nn.Conv1d(256, 256, kernel_size=3, padding=1) # (B, 256, T) -> (B, 256, T)
self.norm2 = LayerNorm(256)
self.linear = nn.Conv1d(256, 1, kernel_size=1) # (B, 256, T) -> (B, 1, T)
self.relu = nn.ReLU()
self.dropout = nn.Dropout(0.1)
def forward(self, x, x_mask):
"""
=====inputs=====
x: (B, 192, T)
x_mask: (B, 1, T)
=====outputs=====
output: (B, 1, T)
"""
x = self.conv1(x * x_mask) # (B, 192, T) -> (B, 256, T)
x = self.relu(x)
x = self.norm1(x)
x = self.dropout(x)
x = self.conv2(x * x_mask) # (B, 256, T) -> (B, 256, T)
x = self.relu(x)
x = self.norm2(x)
x = self.dropout(x)
x = self.linear(x * x_mask) # (B, 256, T) -> (B, 1, T)
output = x * x_mask
return output |