Spaces:
Sleeping
Sleeping
File size: 1,190 Bytes
14d1720 |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
class ResStack(nn.Module):
def __init__(self, channel, dilation=1):
super(ResStack, self).__init__()
self.block = nn.Sequential(
nn.LeakyReLU(0.2),
nn.ReflectionPad1d(dilation),
nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=3, dilation=dilation)),
nn.LeakyReLU(0.2),
nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=1)),
)
self.shortcut = nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=1))
def forward(self, x):
return self.shortcut(x) + self.block(x)
def remove_weight_norm(self):
nn.utils.remove_weight_norm(self.block[2])
nn.utils.remove_weight_norm(self.block[4])
nn.utils.remove_weight_norm(self.shortcut)
# def _remove_weight_norm(m):
# try:
# torch.nn.utils.remove_weight_norm(m)
# except ValueError: # this module didn't have weight norm
# return
#
# self.apply(_remove_weight_norm)
|