def __init__(
self,
n_atom_basis=128,
n_filters=128,
n_interactions=3,
cutoff=5.0,
n_gaussians=25,
normalize_filter=False,
coupled_interactions=False,
return_intermediate=False,
max_z=100,
cutoff_network=HardCutoff,
trainable_gaussians=False,
distance_expansion=None,
charged_systems=False,
):
super(SchNet, self).__init__()
self.n_atom_basis = n_atom_basis
# make a lookup table to store embeddings for each element (up to atomic
# number max_z) each of which is a vector of size n_atom_basis
self.embedding = nn.Embedding(max_z, n_atom_basis, padding_idx=0)
#self.linear1 = nn.Linear(n_atom_basis, 128)
#self.dropout = nn.Dropout(0.05)
# layer for computing interatomic distances
self.distances = AtomDistances()
# layer for expanding interatomic distances in a basis
if distance_expansion is None:
self.distance_expansion = GaussianSmearing(
0.0, cutoff, n_gaussians, trainable=trainable_gaussians)
else:
self.distance_expansion = distance_expansion
# block for computing interaction
if coupled_interactions:
# use the same SchNetInteraction instance (hence the same weights)
self.interactions = nn.ModuleList([
SchNetInteraction(
n_atom_basis=n_atom_basis,
n_spatial_basis=n_gaussians,
n_filters=n_filters,
cutoff_network=cutoff_network,
cutoff=cutoff,
normalize_filter=normalize_filter,
)
] * n_interactions)
else:
# use one SchNetInteraction instance for each interaction
self.interactions = nn.ModuleList([
SchNetInteraction(
n_atom_basis=n_atom_basis,
n_spatial_basis=n_gaussians,
n_filters=n_filters,
cutoff_network=cutoff_network,
cutoff=cutoff,
normalize_filter=normalize_filter,
) for _ in range(n_interactions)
])
# set attributes
self.return_intermediate = return_intermediate
self.charged_systems = charged_systems
if charged_systems:
self.charge = nn.Parameter(torch.Tensor(1, n_atom_basis))
self.charge.data.normal_(0, 1.0 / n_atom_basis**0.5)
def __init__(self):
super(Word2Vec, self).__init__()
# 相当于两次全连接映射
self.W = nn.Parameter(-2 * torch.rand(vocab_size, embedding_size) + 1, requires_grad=True).type(dtype)
self.WT = nn.Parameter(-2 * torch.rand(embedding_size, vocab_size) + 1, requires_grad=True).type(dtype)
def __init__(self, channel):
super().__init__()
self.weight = nn.Parameter(torch.zeros(1, channel, 1, 1))
def __init__(self, args, dictionary, embed_tokens, no_encoder_attn=False):
super().__init__(dictionary)
self.register_buffer('version', torch.Tensor([3]))
self.dropout = args.dropout
self.share_input_output_embed = args.share_decoder_input_output_embed
input_embed_dim = embed_tokens.embedding_dim
embed_dim = args.decoder_embed_dim
self.output_embed_dim = args.decoder_output_dim
self.padding_idx = embed_tokens.padding_idx
self.max_target_positions = args.max_target_positions
self.embed_tokens = embed_tokens
self.embed_scale = math.sqrt(
embed_dim) # todo: try with input_embed_dim
self.project_in_dim = Linear(
input_embed_dim, embed_dim,
bias=False) if embed_dim != input_embed_dim else None
self.embed_positions = PositionalEmbedding(
args.max_target_positions,
embed_dim,
self.padding_idx,
learned=args.decoder_learned_pos,
) if not args.no_token_positional_embeddings else None
self.cross_self_attention = getattr(args, 'cross_self_attention',
False)
self.layer_wise_attention = getattr(args, 'layer_wise_attention',
False)
self.layers = nn.ModuleList([])
self.layers.extend([
TransformerDecoderLayerBN(args, no_encoder_attn)
for _ in range(args.decoder_layers)
])
self.adaptive_softmax = None
self.project_out_dim = Linear(embed_dim, self.output_embed_dim, bias=False) \
if embed_dim != self.output_embed_dim and not args.tie_adaptive_weights else None
if args.adaptive_softmax_cutoff is not None:
self.adaptive_softmax = AdaptiveSoftmax(
len(dictionary),
self.output_embed_dim,
options.eval_str_list(args.adaptive_softmax_cutoff, type=int),
dropout=args.adaptive_softmax_dropout,
adaptive_inputs=embed_tokens
if args.tie_adaptive_weights else None,
factor=args.adaptive_softmax_factor,
tie_proj=args.tie_adaptive_proj,
)
elif not self.share_input_output_embed:
self.embed_out = nn.Parameter(
torch.Tensor(len(dictionary), self.output_embed_dim))
nn.init.normal_(self.embed_out,
mean=0,
std=self.output_embed_dim**-0.5)
if args.decoder_normalize_before and not getattr(
args, 'no_decoder_final_norm', False):
self.layer_norm = LayerNorm(embed_dim)
else:
self.layer_norm = None
def __init__(self, d_hid, eps=1e-3):
super(LayerNorm, self).__init__()
self.eps = eps
self.a_2 = nn.Parameter(torch.ones(d_hid), requires_grad=True)
self.b_2 = nn.Parameter(torch.zeros(d_hid), requires_grad=True)
def test_merge_rows2():
test_size = 100
input_size = 50
output_size = 40
scale = 1.0
w_size = [input_size - 2, output_size]
weight = torch.randn(w_size, dtype=torch.float32)
# weight = torch.arange(0, w_size[0] * w_size[1], dtype=torch.float32).view(w_size)
bias = torch.randn(w_size[1], dtype=weight.dtype)
# bias = None
# Manipulate the weight.
weight = torch.cat((torch.ones([1, output_size], dtype=weight.dtype, device=weight.device), weight), dim=0)
# weight[0][-1] = 1.0
weight = torch.cat((weight, torch.unsqueeze(weight[0] * 1.1, 0)), dim=0)
weight[-1][-1] *= 1.001
# weight = torch.cat((weight, torch.unsqueeze(weight[0] * -0.7, 0)), dim=0)
# print("Weight:")
# print(weight)
proj = weight.new_empty([input_size, input_size])
nn.init.eye_(proj)
# random_input = torch.randn([test_size, input_size], dtype=weight.dtype, device=weight.device)
random_input = torch.ones([test_size, input_size], dtype=weight.dtype, device=weight.device)
from models.ffnn import WorthDense
layer = WorthDense(output_size, False, True, "outer_l2_sum")
_ = layer(random_input)
layer.weight = nn.Parameter(weight, requires_grad=True)
layer.bias = nn.Parameter(bias, requires_grad=True)
layer.in2weight_mat = proj
init_output = layer(random_input)
if bias is not None:
init_cos_sim = pairwise_cos_sim(torch.cat((weight, bias.unsqueeze(0)), dim=0), 0)
else:
init_cos_sim = pairwise_cos_sim(weight.mul(scale), 0)
assert init_cos_sim.shape[0] == output_size
# print("Input:")
# print(random_input)
layer.optimize_weight(0, 0.9999)
print(layer)
new_weight = layer.weight
new_output = layer(random_input)
if bias is not None:
new_cos_sim = pairwise_cos_sim(torch.cat((new_weight, bias.unsqueeze(0)), dim=0), 0)
else:
new_cos_sim = pairwise_cos_sim(new_weight.mul(scale), 0)
assert new_cos_sim.shape[0] == output_size
# print("Init cos sim:")
# print(init_cos_sim)
#
# print("New cos sim:")
# print(new_cos_sim)
#
# print("Gold:")
# print(init_output)
# print("Output:")
# print(new_output)
print(torch.abs(init_output - new_output).sum())
np.testing.assert_almost_equal(init_output.tolist(), new_output.tolist(), decimal=4)
np.testing.assert_almost_equal(init_cos_sim.tolist(), new_cos_sim.tolist(), decimal=4)
def __init__(self, features, eps=1e-6):
super(LayerNorm, self).__init__()
self.a_2 = nn.Parameter(torch.ones(features))
self.b_2 = nn.Parameter(torch.zeros(features))
self.eps = eps
def __init__(self):
super(SegNet, self).__init__()
# initialise network parameters
filter = [64, 128, 256, 512, 512]
self.class_nb = 13
# define encoder decoder layers
self.encoder_block = nn.ModuleList([self.conv_layer([3, filter[0]])])
self.decoder_block = nn.ModuleList([self.conv_layer([filter[0], filter[0]])])
for i in range(4):
self.encoder_block.append(self.conv_layer([filter[i], filter[i + 1]]))
self.decoder_block.append(self.conv_layer([filter[i + 1], filter[i]]))
# define convolution layer
self.conv_block_enc = nn.ModuleList([self.conv_layer([filter[0], filter[0]])])
self.conv_block_dec = nn.ModuleList([self.conv_layer([filter[0], filter[0]])])
for i in range(4):
if i == 0:
self.conv_block_enc.append(self.conv_layer([filter[i + 1], filter[i + 1]]))
self.conv_block_dec.append(self.conv_layer([filter[i], filter[i]]))
else:
self.conv_block_enc.append(nn.Sequential(self.conv_layer([filter[i + 1], filter[i + 1]]),
self.conv_layer([filter[i + 1], filter[i + 1]])))
self.conv_block_dec.append(nn.Sequential(self.conv_layer([filter[i], filter[i]]),
self.conv_layer([filter[i], filter[i]])))
# define task attention layers
self.encoder_att = nn.ModuleList([nn.ModuleList([self.att_layer([filter[0], filter[0], filter[0]])])])
self.decoder_att = nn.ModuleList([nn.ModuleList([self.att_layer([2 * filter[0], filter[0], filter[0]])])])
self.encoder_block_att = nn.ModuleList([self.conv_layer([filter[0], filter[1]])])
self.decoder_block_att = nn.ModuleList([self.conv_layer([filter[0], filter[0]])])
for j in range(3):
if j < 2:
self.encoder_att.append(nn.ModuleList([self.att_layer([filter[0], filter[0], filter[0]])]))
self.decoder_att.append(nn.ModuleList([self.att_layer([2 * filter[0], filter[0], filter[0]])]))
for i in range(4):
self.encoder_att[j].append(self.att_layer([2 * filter[i + 1], filter[i + 1], filter[i + 1]]))
self.decoder_att[j].append(self.att_layer([filter[i + 1] + filter[i], filter[i], filter[i]]))
for i in range(4):
if i < 3:
self.encoder_block_att.append(self.conv_layer([filter[i + 1], filter[i + 2]]))
self.decoder_block_att.append(self.conv_layer([filter[i + 1], filter[i]]))
else:
self.encoder_block_att.append(self.conv_layer([filter[i + 1], filter[i + 1]]))
self.decoder_block_att.append(self.conv_layer([filter[i + 1], filter[i + 1]]))
self.pred_task1 = self.conv_layer([filter[0], self.class_nb], pred=True)
self.pred_task2 = self.conv_layer([filter[0], 1], pred=True)
self.pred_task3 = self.conv_layer([filter[0], 3], pred=True)
# define pooling and unpooling functions
self.down_sampling = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True)
self.up_sampling = nn.MaxUnpool2d(kernel_size=2, stride=2)
self.logsigma = nn.Parameter(torch.FloatTensor([-0.5, -0.5, -0.5]))
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.xavier_normal_(m.weight)
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
nn.init.xavier_normal_(m.weight)
nn.init.constant_(m.bias, 0)
def __init__(self, nfeature): super(Affine, self).__init__() self.nfeature = nfeature self.scale = nn.Parameter(torch.ones(nfeature)) self.bias = nn.Parameter(torch.zeros(nfeature))
def __init__(self,
args,
device,
embedding_matrix,
vocab_size,
word_dict,
embedding_dim,
hidden_size,
dropout=0.2):
super(Encoder, self).__init__()
self.args = args
self.word_dict = word_dict
idx2word = {v: k for k, v in word_dict.items()}
glove_path = 'data/glove_word2vec_300d.txt'
if self.args.glove:
# self.embed = None
self.embed = nn.Embedding.from_pretrained(
get_pretrained_glove(path=glove_path,
idx2word=idx2word,
n_special=len(token)),
freeze=False,
padding_idx=0) # specials: pad, unk, naf_h/t
else:
self.embed = nn.Embedding(vocab_size, embedding_dim)
# self.pmr_init_matrix = torch.randn(args.char_num*args.pmr_size, args.hidden_size*2).to(device) # (6*32, 2h) randomly initialize a pmr matrix
# 2-17
# self.pmr_init_matrix = torch.randn(args.pmr_size, args.hidden_size*2).to(device) # (32, 2h) randomly initialize a pmr matrix
# 2-28
self.pmr_init_matrix = nn.Parameter(
torch.randn(args.pmr_size, args.hidden_size * 2),
requires_grad=True) # (32, 2h) randomly initialize a pmr matrix
self.Wh1 = nn.Linear(hidden_size * 1, hidden_size * 2)
self.Wc1 = nn.Linear(hidden_size * 1, hidden_size * 2)
self.bilstm = nn.LSTM(embedding_dim,
hidden_size * 2,
2,
batch_first=True,
bidirectional=True)
self.bilstm1 = nn.LSTM(embedding_dim,
hidden_size * 2,
2,
batch_first=True,
bidirectional=True)
self.dropout_x = nn.Dropout(dropout)
self.dropout_c = nn.Dropout(dropout)
self.fc = nn.Linear(hidden_size * 4, hidden_size)
self.fc_hid = nn.Linear(hidden_size * 4, hidden_size)
self.fc_hid1 = nn.Linear(hidden_size * 4, hidden_size)
# self.fc_hid_lstm = nn.Linear(hidden_size * 2, hidden_size)
# self.fc_hid_lstm1 = nn.Linear(hidden_size * 2, hidden_size)
self.fc_c = nn.Linear(hidden_size * 4, hidden_size)
self.fc_c1 = nn.Linear(hidden_size * 4, hidden_size)
# self.fc_c_lstm = nn.Linear(hidden_size * 2, hidden_size)
# self.fc_c_lstm1 = nn.Linear(hidden_size * 2, hidden_size)
self.enc_out = nn.Linear(hidden_size * 4, hidden_size * 2)
# self.enc_out_lstm = nn.Linear(hidden_size * 2, hidden_size * 2)
# self.enc_out_lstm1 = nn.Linear(hidden_size * 2, hidden_size * 2)
self.W_s1 = nn.Linear(2 * hidden_size, 350)
self.W_s2 = nn.Linear(350, 30)
self.fc_layer = nn.Linear(30 * 2 * hidden_size, hidden_size * 2)
def __init__(self, num_channels):
super().__init__()
# self.num_channels = num_channels
# register and initialize the weight and bias
self.w = nn.Parameter(1 + 0.1 * torch.randn(num_channels))
self.b = nn.Parameter(0 + 0.1 * torch.randn(num_channels))
def __init__(self, out_size, init):
super(Gauss_NoiseModel, self).__init__() # always call parent's init
self.sigma = nn.Parameter(torch.Tensor(out_size).fill_(init), requires_grad=True)
def __init__(self, out_sizeStudentT, init_scale):
super(StudentT_NoiseModel, self).__init__() # always call parent's init
self.scale = nn.Parameter(torch.Tensor(out_sizeStudentT).fill_(init_scale), requires_grad=True)
def __init__(self, in_channels, out_channels, inter_channels, num_subset=3, num_node=25, num_frame=32,
kernel_size=1, stride=1, glo_reg_s=True, att_s=True, glo_reg_t=True, att_t=True,
use_temporal_att=True, use_spatial_att=True, attentiondrop=0, use_pes=True, use_pet=True):
super(STAttentionBlock, self).__init__()
self.inter_channels = inter_channels
self.out_channels = out_channels
self.in_channels = in_channels
self.num_subset = num_subset
self.glo_reg_s = glo_reg_s
self.att_s = att_s
self.glo_reg_t = glo_reg_t
self.att_t = att_t
self.use_pes = use_pes
self.use_pet = use_pet
backward_mask = torch.triu(torch.ones(num_frame, num_frame))
self.register_buffer('backward_mask', backward_mask)
pad = int((kernel_size - 1) / 2)
self.use_spatial_att = use_spatial_att
if use_spatial_att:
atts = torch.zeros((1, num_subset, num_node, num_node))
self.register_buffer('atts', atts)
self.pes = PositionalEncoding(in_channels, num_node, num_frame, 'spatial')
self.ff_nets = nn.Sequential(
nn.Conv2d(out_channels, out_channels, 1, 1, padding=0, bias=True),
nn.BatchNorm2d(out_channels),
)
if att_s:
self.in_nets = nn.Conv2d(in_channels, 2 * num_subset * inter_channels, 1, bias=True)
self.alphas = nn.Parameter(torch.ones(1, num_subset, 1, 1), requires_grad=True)
if glo_reg_s:
self.attention0s = nn.Parameter(torch.ones(1, num_subset, num_node, num_node) / num_node,
requires_grad=True)
self.out_nets = nn.Sequential(
nn.Conv2d(in_channels * num_subset, out_channels, 1, bias=True),
nn.BatchNorm2d(out_channels),
)
else:
self.out_nets = nn.Sequential(
nn.Conv2d(in_channels, out_channels, (1, 3), padding=(0, 1), bias=True, stride=1),
nn.BatchNorm2d(out_channels),
)
self.use_temporal_att = use_temporal_att
if use_temporal_att:
attt = torch.zeros((1, num_subset, num_frame, num_frame))
self.register_buffer('attt', attt)
self.pet = PositionalEncoding(out_channels, num_node, num_frame, 'temporal')
self.ff_nett = nn.Sequential(
nn.Conv2d(out_channels, out_channels, (kernel_size, 1), (stride, 1), padding=(pad, 0), bias=True),
nn.BatchNorm2d(out_channels),
)
if att_t:
self.in_nett = nn.Conv2d(out_channels, 4 * num_subset * inter_channels, 1, bias=True)
# self.alphat = nn.Parameter(torch.ones(1, num_subset, 1, 1), requires_grad=True)
self.alphat_0 = nn.Parameter(torch.ones(1, num_subset, 1, 1), requires_grad=True)
self.alphat_1 = nn.Parameter(torch.ones(1, num_subset, 1, 1), requires_grad=True)
if glo_reg_t:
self.attention0t = nn.Parameter(torch.zeros(1, num_subset, num_frame, num_frame) + torch.eye(num_frame),
requires_grad=True)
self.out_nett = nn.Sequential(
nn.Conv2d(out_channels * num_subset * 2, out_channels, 1, bias=True),
nn.BatchNorm2d(out_channels),
)
else:
self.out_nett = nn.Sequential(
nn.Conv2d(out_channels, out_channels, (7, 1), padding=(3, 0), bias=True, stride=(stride, 1)),
nn.BatchNorm2d(out_channels),
)
self.out_nett_extend = nn.Sequential(nn.Conv2d(out_channels, out_channels, (7, 1), padding=(3, 0), bias=True, stride=(stride, 1)),nn.BatchNorm2d(out_channels),)
if in_channels != out_channels or stride != 1:
if use_spatial_att:
self.downs1 = nn.Sequential(
nn.Conv2d(in_channels, out_channels, 1, bias=True),
nn.BatchNorm2d(out_channels),
)
self.downs2 = nn.Sequential(
nn.Conv2d(in_channels, out_channels, 1, bias=True),
nn.BatchNorm2d(out_channels),
)
if use_temporal_att:
self.downt1 = nn.Sequential(
nn.Conv2d(out_channels, out_channels, 1, 1, bias=True),
nn.BatchNorm2d(out_channels),
)
self.downt2 = nn.Sequential(
nn.Conv2d(out_channels, out_channels, (kernel_size, 1), (stride, 1), padding=(pad, 0), bias=True),
nn.BatchNorm2d(out_channels),
)
self.downt3 = nn.Sequential(
nn.Conv2d(out_channels, out_channels, (kernel_size, 1), (stride, 1), padding=(pad, 0), bias=True),
nn.BatchNorm2d(out_channels),
)
else:
if use_spatial_att:
self.downs1 = lambda x: x
self.downs2 = lambda x: x
if use_temporal_att:
self.downt1 = lambda x: x
self.downt2 = lambda x: x
self.downt3 = lambda x: x
self.soft = nn.Softmax(-2)
self.tan = nn.Tanh()
self.relu = nn.LeakyReLU(0.1)
self.elu = nn.ELU()
self.drop = nn.Dropout(attentiondrop)
def __init__(self, init_value=1.0):
super(Scale, self).__init__()
self.scale = nn.Parameter(torch.FloatTensor([init_value]))
def __init__(self):
super().__init__()
self.p = nn.Parameter(torch.ones(2, 2))
def test_merge_rows():
test_size = 100
input_size = 500
output_size = 400
scale = 1.0
w_size = [input_size - 2, output_size]
weight = torch.randn(w_size, dtype=torch.float32)
# weight = torch.arange(0, w_size[0] * w_size[1], dtype=torch.float32).view(w_size)
bias = torch.randn(w_size[1], dtype=weight.dtype)
# bias = None
# Manipulate the weight.
weight = torch.cat((weight, torch.unsqueeze(weight[0] * 1.7, 0)), dim=0)
weight = torch.cat((weight, torch.unsqueeze(weight[0] * -0.7, 0)), dim=0)
print("Weight:")
print(weight)
proj = weight.new_empty([input_size, input_size])
nn.init.eye_(proj)
random_input = torch.randn([test_size, input_size], dtype=weight.dtype, device=weight.device)
# random_input = torch.ones([test_size, input_size], dtype=weight.dtype, device=weight.device)
from models.ffnn import WorthDense
layer = WorthDense(output_size, False, True, "outer_l2_sum")
_ = layer(random_input)
layer.weight = nn.Parameter(weight, requires_grad=True)
layer.bias = nn.Parameter(bias, requires_grad=True)
init_output = layer(random_input)
# init_output = nn.functional.linear(random_input, weight.mul(scale).t(), bias)
if bias is not None:
init_cos_sim = pairwise_cos_sim(torch.cat((weight, bias.unsqueeze(0)), dim=0), 0)
else:
init_cos_sim = pairwise_cos_sim(weight.mul(scale), 0)
init_other_cos_sim = pairwise_cos_sim(weight.mul(scale)[:-2], 1)
assert init_cos_sim.shape[0] == output_size
new_weight, new_trans_mat = merge_rows(weight, proj, 0, input_size - 1)
new_weight, new_trans_mat = merge_rows(new_weight, new_trans_mat, 0, input_size - 2)
print("Proj:")
print(new_trans_mat)
print("Input:")
print(random_input)
proj_input = nn.functional.linear(random_input, new_trans_mat.t(), None)
print("Proj input:")
print(proj_input)
print("Proj weight:")
print(new_weight)
layer.weight = nn.Parameter(new_weight, requires_grad=True)
layer.in2weight_mat = new_trans_mat
new_output = layer(random_input)
# new_output = nn.functional.linear(proj_input, new_weight.mul(scale).t(), bias)
if bias is not None:
new_cos_sim = pairwise_cos_sim(torch.cat((new_weight, bias.unsqueeze(0)), dim=0), 0)
else:
new_cos_sim = pairwise_cos_sim(new_weight.mul(scale), 0)
new_other_cos_sim = pairwise_cos_sim(new_weight.mul(scale), 1)
assert new_cos_sim.shape[0] == output_size
print("Init cos sim:")
print(init_cos_sim)
print("New cos sim:")
print(new_cos_sim)
print("Gold:")
print(init_output)
print("Output:")
print(new_output)
np.testing.assert_almost_equal(init_output.tolist(), new_output.tolist(), decimal=5)
np.testing.assert_almost_equal(init_cos_sim.tolist(), new_cos_sim.tolist(), decimal=5)
np.testing.assert_almost_equal(init_other_cos_sim.tolist(), new_other_cos_sim.tolist(), decimal=7)
def __init__(self, args):
KalmanBasis.__init__(self, 6, 2, args)
self.dt = args.dt
self._position_std_x = nn.Parameter(torch.ones(1) * 5)
self._position_std_y = nn.Parameter(torch.ones(1) * 5)
self._velocity_std_x = nn.Parameter(torch.ones(1) * 10)
self._velocity_std_y = nn.Parameter(torch.ones(1) * 10)
self._acceleration_std_x = nn.Parameter(torch.ones(1) * 5)
self._acceleration_std_y = nn.Parameter(torch.ones(1) * 5)
self._jerk_std_x = nn.Parameter(torch.ones(1) * 3)
self._jerk_std_y = nn.Parameter(torch.ones(1) * 3)
self._n_command = 2
coef_G = torch.randn(self._state_size, self._state_size)
self._coef_G = nn.Parameter(coef_G)
GR = torch.randn(2, 2)
self._GR = nn.Parameter(GR)
self.Id = nn.Parameter(torch.eye(self._state_size), requires_grad=False)
dt = args.dt
# Transition matrix that defines evolution of position over a time step for a given state
self._F = nn.Parameter(torch.eye(self._state_size), requires_grad=False)
self._F[0, 2] = dt
self._F[0, 4] = dt * dt / 2
self._F[1, 3] = dt
self._F[1, 5] = dt * dt / 2
self._F[2, 4] = dt
self._F[3, 5] = dt
# Command matrix that defines how commands modify the state
self._B = nn.Parameter(torch.zeros(self._state_size, self._n_command), requires_grad=False) # actions are accelerations
self._B[0, 0] = 0#dt * dt * dt / 6
self._B[1, 1] = 0#dt * dt * dt / 6
self._B[2, 0] = 0#dt * dt / 2
self._B[3, 1] = 0#dt * dt / 2
self._B[4, 0] = dt
self._B[5, 1] = dt
def __init__(
self,
in_channels,
out_channels,
kernel_size,
sample_rate=16000,
min_low_hz=50,
min_band_hz=50,
stride=1,
padding=0,
dilation=1,
bias=False,
groups=1,
):
super().__init__()
if in_channels != 1:
msg = (f"SincConv1d only supports one input channel. "
f"Here, in_channels = {in_channels}.")
raise ValueError(msg)
self.in_channels = in_channels
self.out_channels = out_channels
if kernel_size % 2 == 0:
msg = (f"SincConv1d only support odd kernel size. "
f"Here, kernel_size = {kernel_size}.")
raise ValueError(msg)
self.kernel_size = kernel_size
self.stride = stride
self.padding = padding
self.dilation = dilation
if bias:
raise ValueError("SincConv1d does not support bias.")
if groups > 1:
raise ValueError("SincConv1d does not support groups.")
self.sample_rate = sample_rate
self.min_low_hz = min_low_hz
self.min_band_hz = min_band_hz
# initialize filterbanks such that they are equally spaced in Mel scale
low_hz = 30
high_hz = self.sample_rate / 2 - (self.min_low_hz + self.min_band_hz)
mel = np.linspace(self.to_mel(low_hz), self.to_mel(high_hz),
self.out_channels + 1)
hz = self.to_hz(mel)
# filter lower frequency (out_channels, 1)
self.low_hz_ = nn.Parameter(torch.Tensor(hz[:-1]).view(-1, 1))
# filter frequency band (out_channels, 1)
self.band_hz_ = nn.Parameter(torch.Tensor(np.diff(hz)).view(-1, 1))
# Half Hamming half window
n_lin = torch.linspace(0,
self.kernel_size / 2 - 1,
steps=int((self.kernel_size / 2)))
self.window_ = 0.54 - 0.46 * torch.cos(
2 * math.pi * n_lin / self.kernel_size)
# (kernel_size, 1)
# Due to symmetry, I only need half of the time axes
n = (self.kernel_size - 1) / 2.0
self.n_ = 2 * math.pi * torch.arange(-n, 0).view(1,
-1) / self.sample_rate
def __init__(self, win_len, stride, F_size):
super(ConcentrationLoss, self).__init__()
self.win_len = win_len
self.grid = nn.Parameter(create_grid(F_size), requires_grad=False)
self.F_size = F_size
self.stride = stride
def __init__(self, n_hidden):
super(Discriminator, self).__init__()
self.weight = nn.Parameter(
nn.init.xavier_uniform_(torch.empty(n_hidden, n_hidden)))
def __init__(self, num_features, eps=1e-5):
super().__init__()
self.eps = eps
self.num_features = num_features
self.gamma = nn.Parameter(torch.ones(self.num_features))
self.beta = nn.Parameter(torch.zeros(self.num_features))
def __init__(self):
super(MyDense, self).__init__()
self.params = nn.ParameterList([])
self.params.append(nn.Parameter(torch.rand((36 ,4), requires_grad=True)))
def __init__(self, model, loss_fn):
super(MultiTask3Loss, self).__init__()
self.model = model
self.loss_fn = loss_fn
self.eta = nn.Parameter(torch.FloatTensor(len(loss_fn)).fill_(1.), requires_grad=True)
def __init__(self, input_size, output_size):
super(NonBiasLinear, self).__init__()
self.weight = nn.Parameter(data=torch.randn(input_size, output_size), requires_grad=True)
def __init__(self, n_categories, batch_shape: Size = torch.Size()):
super().__init__(batch_shape=torch.Size(batch_shape))
self.logits = nn.Parameter(torch.randn(*batch_shape, n_categories))
def __init__(self, drop_rate):
super(_VDropout, self).__init__()
self.drop_rate = drop_rate
self.dummy_param = nn.Parameter(torch.empty(0))
def from_distribution(dist: Categorical):
new = CategoricalParams.__new__(CategoricalParams)
super(CategoricalParams, new).__init__(batch_shape=dist.batch_shape)
new.logits = nn.Parameter(dist.logits)
return new
def __init__(self, channel, size=4):
super().__init__()
self.input = nn.Parameter(torch.randn(1, channel, size, size))
def forward(self, inputs):
# convert inputs from BCHW -> BHWC
inputs = inputs.permute(0, 2, 3, 1).contiguous()
input_shape = inputs.shape
batch_size, height, width, channels = input_shape
# Flatten input
flat_input = inputs.view(-1, self._embedding_dim)
# Calculate distances
distances = (
torch.sum(flat_input ** 2, dim=1, keepdim=True)
+ torch.sum(self._embedding.weight ** 2, dim=1)
- 2 * torch.matmul(flat_input, self._embedding.weight.t())
)
# Encoding
encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)
encodings = torch.zeros(
encoding_indices.shape[0], self._num_embeddings, device=inputs.device
)
encodings.scatter_(1, encoding_indices, 1)
# Quantize and unflatten
quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape)
# Use EMA to update the embedding vectors
if self.training:
self._ema_cluster_size = self._ema_cluster_size * self._decay + (
1 - self._decay
) * torch.sum(encodings, 0)
# Laplace smoothing of the cluster size
n = torch.sum(self._ema_cluster_size.data)
self._ema_cluster_size = (
(self._ema_cluster_size + self._epsilon)
/ (n + self._num_embeddings * self._epsilon)
* n
)
dw = torch.matmul(encodings.t(), flat_input)
self._ema_w = nn.Parameter(
self._ema_w * self._decay + (1 - self._decay) * dw
)
self._embedding.weight = nn.Parameter(
self._ema_w / self._ema_cluster_size.unsqueeze(1)
)
# Loss
e_latent_loss = F.mse_loss(quantized.detach(), inputs)
loss = self._commitment_cost * e_latent_loss
# Straight Through Estimator
quantized = inputs + (quantized - inputs).detach()
avg_probs = torch.mean(encodings, dim=0)
perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
encoding_indices = encoding_indices.view(batch_size, height, width)
# convert quantized from BHWC -> BCHW
return (
loss,
quantized.permute(0, 3, 1, 2).contiguous(),
perplexity,
encoding_indices,
)
