Spaces:
Build error
Build error
| import torch | |
| import torch.nn as nn | |
| from torch.distributions import MultivariateNormal | |
| import math | |
| import numpy as np | |
| from helper import gaussian_2d | |
| from config.GlobalVariables import * | |
| class SynthesisNetwork(nn.Module): | |
| def __init__(self, weight_dim=512, num_layers=3, scale_sd=1, clamp_mdn=0, sentence_loss=True, word_loss=True, segment_loss=True, TYPE_A=True, TYPE_B=True, TYPE_C=True, TYPE_D=True, ORIGINAL=True, REC=True): | |
| super(SynthesisNetwork, self).__init__() | |
| self.num_mixtures = 20 | |
| self.num_layers = num_layers | |
| self.weight_dim = weight_dim | |
| self.device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| self.sentence_loss = sentence_loss | |
| self.word_loss = word_loss | |
| self.segment_loss = segment_loss | |
| self.ORIGINAL = ORIGINAL | |
| self.TYPE_A = TYPE_A | |
| self.TYPE_B = TYPE_B | |
| self.TYPE_C = TYPE_C | |
| self.TYPE_D = TYPE_D | |
| self.REC = REC | |
| self.magic_lstm = nn.LSTM(self.weight_dim, self.weight_dim, batch_first=True, num_layers=self.num_layers) | |
| self.char_vec_fc_1 = nn.Linear(len(CHARACTERS), self.weight_dim) | |
| self.char_vec_relu_1 = nn.LeakyReLU(negative_slope=0.1) | |
| self.char_lstm_1 = nn.LSTM(self.weight_dim, self.weight_dim, batch_first=True, num_layers=self.num_layers) | |
| self.char_vec_fc2_1 = nn.Linear(self.weight_dim, self.weight_dim * self.weight_dim) | |
| # inference | |
| self.inf_state_fc1 = nn.Linear(3, self.weight_dim) | |
| self.inf_state_relu = nn.LeakyReLU(negative_slope=0.1) | |
| self.inf_state_lstm = nn.LSTM(self.weight_dim, self.weight_dim, batch_first=True, num_layers=self.num_layers) | |
| self.W_lstm = nn.LSTM(self.weight_dim, self.weight_dim, batch_first=True, num_layers=self.num_layers) | |
| # generation | |
| self.gen_state_fc1 = nn.Linear(3, self.weight_dim) | |
| self.gen_state_relu = nn.LeakyReLU(negative_slope=0.1) | |
| self.gen_state_lstm1 = nn.LSTM(self.weight_dim, self.weight_dim, batch_first=True, num_layers=self.num_layers) | |
| self.gen_state_lstm2 = nn.LSTM(self.weight_dim * 2, self.weight_dim * 2, batch_first=True, num_layers=self.num_layers) | |
| self.gen_state_fc2 = nn.Linear(self.weight_dim * 2, self.num_mixtures * 6 + 1) | |
| self.term_fc1 = nn.Linear(self.weight_dim * 2, self.weight_dim) | |
| self.term_relu1 = nn.LeakyReLU(negative_slope=0.1) | |
| self.term_fc2 = nn.Linear(self.weight_dim, self.weight_dim) | |
| self.term_relu2 = nn.LeakyReLU(negative_slope=0.1) | |
| self.term_fc3 = nn.Linear(self.weight_dim, 1) | |
| self.term_sigmoid = nn.Sigmoid() | |
| self.mdn_sigmoid = nn.Sigmoid() | |
| self.mdn_tanh = nn.Tanh() | |
| self.mdn_softmax = nn.Softmax(dim=1) | |
| self.scale_sd = scale_sd # how much to scale the standard deviation of the gaussians | |
| self.clamp_mdn = clamp_mdn # total percent of disrubution to allow sampling from | |
| self.mdn_bce_loss = nn.BCEWithLogitsLoss() | |
| self.term_bce_loss = nn.BCEWithLogitsLoss() | |
| def forward(self, inputs): | |
| [sentence_level_stroke_in, sentence_level_stroke_out, sentence_level_stroke_length, sentence_level_term, sentence_level_char, sentence_level_char_length, word_level_stroke_in, word_level_stroke_out, word_level_stroke_length, word_level_term, word_level_char, word_level_char_length, segment_level_stroke_in, segment_level_stroke_out, segment_level_stroke_length, segment_level_term, segment_level_char, segment_level_char_length] = inputs | |
| ALL_sentence_W_consistency_loss = [] | |
| ALL_ORIGINAL_sentence_termination_loss = [] | |
| ALL_ORIGINAL_sentence_loc_reconstruct_loss = [] | |
| ALL_ORIGINAL_sentence_touch_reconstruct_loss = [] | |
| ALL_TYPE_A_sentence_termination_loss = [] | |
| ALL_TYPE_A_sentence_loc_reconstruct_loss = [] | |
| ALL_TYPE_A_sentence_touch_reconstruct_loss = [] | |
| ALL_TYPE_A_sentence_WC_reconstruct_loss = [] | |
| ALL_TYPE_B_sentence_termination_loss = [] | |
| ALL_TYPE_B_sentence_loc_reconstruct_loss = [] | |
| ALL_TYPE_B_sentence_touch_reconstruct_loss = [] | |
| ALL_TYPE_B_sentence_WC_reconstruct_loss = [] | |
| ALL_word_W_consistency_loss = [] | |
| ALL_ORIGINAL_word_termination_loss = [] | |
| ALL_ORIGINAL_word_loc_reconstruct_loss = [] | |
| ALL_ORIGINAL_word_touch_reconstruct_loss = [] | |
| ALL_TYPE_A_word_termination_loss = [] | |
| ALL_TYPE_A_word_loc_reconstruct_loss = [] | |
| ALL_TYPE_A_word_touch_reconstruct_loss = [] | |
| ALL_TYPE_A_word_WC_reconstruct_loss = [] | |
| ALL_TYPE_B_word_termination_loss = [] | |
| ALL_TYPE_B_word_loc_reconstruct_loss = [] | |
| ALL_TYPE_B_word_touch_reconstruct_loss = [] | |
| ALL_TYPE_B_word_WC_reconstruct_loss = [] | |
| ALL_TYPE_C_word_termination_loss = [] | |
| ALL_TYPE_C_word_loc_reconstruct_loss = [] | |
| ALL_TYPE_C_word_touch_reconstruct_loss = [] | |
| ALL_TYPE_C_word_WC_reconstruct_loss = [] | |
| ALL_TYPE_D_word_termination_loss = [] | |
| ALL_TYPE_D_word_loc_reconstruct_loss = [] | |
| ALL_TYPE_D_word_touch_reconstruct_loss = [] | |
| ALL_TYPE_D_word_WC_reconstruct_loss = [] | |
| ALL_word_Wcs_reconstruct_TYPE_A = [] | |
| ALL_word_Wcs_reconstruct_TYPE_B = [] | |
| ALL_word_Wcs_reconstruct_TYPE_C = [] | |
| ALL_word_Wcs_reconstruct_TYPE_D = [] | |
| SUPER_ALL_segment_W_consistency_loss = [] | |
| SUPER_ALL_ORIGINAL_segment_termination_loss = [] | |
| SUPER_ALL_ORIGINAL_segment_loc_reconstruct_loss = [] | |
| SUPER_ALL_ORIGINAL_segment_touch_reconstruct_loss = [] | |
| SUPER_ALL_TYPE_A_segment_termination_loss = [] | |
| SUPER_ALL_TYPE_A_segment_loc_reconstruct_loss = [] | |
| SUPER_ALL_TYPE_A_segment_touch_reconstruct_loss = [] | |
| SUPER_ALL_TYPE_A_segment_WC_reconstruct_loss = [] | |
| SUPER_ALL_TYPE_B_segment_termination_loss = [] | |
| SUPER_ALL_TYPE_B_segment_loc_reconstruct_loss = [] | |
| SUPER_ALL_TYPE_B_segment_touch_reconstruct_loss = [] | |
| SUPER_ALL_TYPE_B_segment_WC_reconstruct_loss = [] | |
| SUPER_ALL_segment_Wcs_reconstruct_TYPE_A = [] | |
| SUPER_ALL_segment_Wcs_reconstruct_TYPE_B = [] | |
| # if self.sentece_loss: | |
| for uid in range(len(sentence_level_stroke_in)): | |
| if self.sentence_loss: | |
| user_sentence_level_stroke_in = sentence_level_stroke_in[uid] | |
| user_sentence_level_stroke_out = sentence_level_stroke_out[uid] | |
| user_sentence_level_stroke_length = sentence_level_stroke_length[uid] | |
| user_sentence_level_term = sentence_level_term[uid] | |
| user_sentence_level_char = sentence_level_char[uid] | |
| user_sentence_level_char_length = sentence_level_char_length[uid] | |
| sentence_batch_size = len(user_sentence_level_stroke_in) | |
| sentence_inf_state_out = self.inf_state_fc1(user_sentence_level_stroke_out) | |
| sentence_inf_state_out = self.inf_state_relu(sentence_inf_state_out) | |
| sentence_inf_state_out, (c,h) = self.inf_state_lstm(sentence_inf_state_out) | |
| sentence_gen_state_out = self.gen_state_fc1(user_sentence_level_stroke_in) | |
| sentence_gen_state_out = self.gen_state_relu(sentence_gen_state_out) | |
| sentence_gen_state_out, (c,h) = self.gen_state_lstm1(sentence_gen_state_out) | |
| sentence_Ws = [] | |
| sentence_Wc_rec_TYPE_ = [] | |
| sentence_SPLITS = [] | |
| sentence_Cs_1 = [] | |
| sentence_unique_char_matrices_1 = [] | |
| for sentence_batch_id in range(sentence_batch_size): | |
| curr_seq_len = user_sentence_level_stroke_length[sentence_batch_id][0] | |
| curr_char_len = user_sentence_level_char_length[sentence_batch_id][0] | |
| char_vector = torch.eye(len(CHARACTERS))[user_sentence_level_char[sentence_batch_id][:curr_char_len]].to(self.device) | |
| current_term = user_sentence_level_term[sentence_batch_id][:curr_seq_len].unsqueeze(-1) | |
| split_ids = torch.nonzero(current_term)[:,0] | |
| char_vector_1 = self.char_vec_fc_1(char_vector) | |
| char_vector_1 = self.char_vec_relu_1(char_vector_1) | |
| unique_char_matrices_1 = [] | |
| for cid in range(len(char_vector)): | |
| # Tower 1 | |
| unique_char_vector_1 = char_vector_1[cid:cid+1] | |
| unique_char_input_1 = unique_char_vector_1.unsqueeze(0) | |
| unique_char_out_1, (c,h) = self.char_lstm_1(unique_char_input_1) | |
| unique_char_out_1 = unique_char_out_1.squeeze(0) | |
| unique_char_out_1 = self.char_vec_fc2_1(unique_char_out_1) | |
| unique_char_matrix_1 = unique_char_out_1.view([-1,1,self.weight_dim,self.weight_dim]) | |
| unique_char_matrix_1 = unique_char_matrix_1.squeeze(1) | |
| unique_char_matrices_1.append(unique_char_matrix_1) | |
| # Tower 1 | |
| char_out_1 = char_vector_1.unsqueeze(0) | |
| char_out_1, (c,h) = self.char_lstm_1(char_out_1) | |
| char_out_1 = char_out_1.squeeze(0) | |
| char_out_1 = self.char_vec_fc2_1(char_out_1) | |
| char_matrix_1 = char_out_1.view([-1,1,self.weight_dim,self.weight_dim]) | |
| char_matrix_1 = char_matrix_1.squeeze(1) | |
| char_matrix_inv_1 = torch.inverse(char_matrix_1) | |
| W_c_t = sentence_inf_state_out[sentence_batch_id][:curr_seq_len] | |
| W_c = torch.stack([W_c_t[i] for i in split_ids]) | |
| # W = torch.bmm(char_matrix_inv, W_c.unsqueeze(2)).squeeze(-1) | |
| # C1C2C3W = Wc | |
| # W = C3-1 C2-1 C1-1 Wc | |
| W = torch.bmm(char_matrix_inv_1, | |
| W_c.unsqueeze(2)).squeeze(-1) | |
| sentence_Ws.append(W) | |
| sentence_Wc_rec_TYPE_.append(W_c) | |
| sentence_Cs_1.append(char_matrix_1) | |
| sentence_SPLITS.append(split_ids) | |
| sentence_unique_char_matrices_1.append(unique_char_matrices_1) | |
| sentence_Ws_stacked = torch.cat(sentence_Ws, 0) | |
| sentence_Ws_reshaped = sentence_Ws_stacked.view([-1,self.weight_dim]) | |
| sentence_W_mean = sentence_Ws_reshaped.mean(0) | |
| sentence_W_mean_repeat = sentence_W_mean.repeat(sentence_Ws_reshaped.size(0),1) | |
| sentence_Ws_consistency_loss = torch.mean(torch.mean(torch.mul(sentence_W_mean_repeat - sentence_Ws_reshaped, sentence_W_mean_repeat - sentence_Ws_reshaped), -1)) | |
| ALL_sentence_W_consistency_loss.append(sentence_Ws_consistency_loss) | |
| ORIGINAL_sentence_termination_loss = [] | |
| ORIGINAL_sentence_loc_reconstruct_loss = [] | |
| ORIGINAL_sentence_touch_reconstruct_loss = [] | |
| TYPE_A_sentence_termination_loss = [] | |
| TYPE_A_sentence_loc_reconstruct_loss = [] | |
| TYPE_A_sentence_touch_reconstruct_loss = [] | |
| TYPE_B_sentence_termination_loss = [] | |
| TYPE_B_sentence_loc_reconstruct_loss = [] | |
| TYPE_B_sentence_touch_reconstruct_loss = [] | |
| sentence_Wcs_reconstruct_TYPE_A = [] | |
| sentence_Wcs_reconstruct_TYPE_B = [] | |
| for sentence_batch_id in range(sentence_batch_size): | |
| sentence_level_gen_encoded = sentence_gen_state_out[sentence_batch_id][:user_sentence_level_stroke_length[sentence_batch_id][0]] | |
| sentence_level_target_eos = user_sentence_level_stroke_out[sentence_batch_id][:user_sentence_level_stroke_length[sentence_batch_id][0]][:,2] | |
| sentence_level_target_x = user_sentence_level_stroke_out[sentence_batch_id][:user_sentence_level_stroke_length[sentence_batch_id][0]][:,0:1] | |
| sentence_level_target_y = user_sentence_level_stroke_out[sentence_batch_id][:user_sentence_level_stroke_length[sentence_batch_id][0]][:,1:2] | |
| sentence_level_target_term = user_sentence_level_term[sentence_batch_id][:user_sentence_level_stroke_length[sentence_batch_id][0]] | |
| # ORIGINAL | |
| if self.ORIGINAL: | |
| sentence_W_lstm_in_ORIGINAL = [] | |
| curr_id = 0 | |
| for i in range(user_sentence_level_stroke_length[sentence_batch_id][0]): | |
| sentence_W_lstm_in_ORIGINAL.append(sentence_Wc_rec_TYPE_[sentence_batch_id][curr_id]) | |
| if i in sentence_SPLITS[sentence_batch_id]: | |
| curr_id += 1 | |
| sentence_W_lstm_in_ORIGINAL = torch.stack(sentence_W_lstm_in_ORIGINAL) | |
| sentence_Wc_t_ORIGINAL = sentence_W_lstm_in_ORIGINAL | |
| sentence_gen_lstm2_in_ORIGINAL = torch.cat([sentence_level_gen_encoded, sentence_Wc_t_ORIGINAL], -1) | |
| sentence_gen_lstm2_in_ORIGINAL = sentence_gen_lstm2_in_ORIGINAL.unsqueeze(0) | |
| sentence_gen_out_ORIGINAL,(c,h) = self.gen_state_lstm2(sentence_gen_lstm2_in_ORIGINAL) | |
| sentence_gen_out_ORIGINAL = sentence_gen_out_ORIGINAL.squeeze(0) | |
| mdn_out_ORIGINAL = self.gen_state_fc2(sentence_gen_out_ORIGINAL) | |
| eos_ORIGINAL = mdn_out_ORIGINAL[:,0:1] | |
| [mu1_ORIGINAL, mu2_ORIGINAL, sig1_ORIGINAL, sig2_ORIGINAL, rho_ORIGINAL, pi_ORIGINAL] = torch.split(mdn_out_ORIGINAL[:,1:], self.num_mixtures, 1) | |
| sig1_ORIGINAL = sig1_ORIGINAL.exp() + 1e-3 | |
| sig2_ORIGINAL = sig2_ORIGINAL.exp() + 1e-3 | |
| rho_ORIGINAL = self.mdn_tanh(rho_ORIGINAL) | |
| pi_ORIGINAL = self.mdn_softmax(pi_ORIGINAL) | |
| term_out_ORIGINAL = self.term_fc1(sentence_gen_out_ORIGINAL) | |
| term_out_ORIGINAL = self.term_relu1(term_out_ORIGINAL) | |
| term_out_ORIGINAL = self.term_fc2(term_out_ORIGINAL) | |
| term_out_ORIGINAL = self.term_relu2(term_out_ORIGINAL) | |
| term_out_ORIGINAL = self.term_fc3(term_out_ORIGINAL) | |
| term_pred_ORIGINAL = self.term_sigmoid(term_out_ORIGINAL) | |
| gaussian_ORIGINAL = gaussian_2d(sentence_level_target_x, sentence_level_target_y, mu1_ORIGINAL, mu2_ORIGINAL, sig1_ORIGINAL, sig2_ORIGINAL, rho_ORIGINAL) | |
| loss_gaussian_ORIGINAL = - torch.log(torch.sum(pi_ORIGINAL*gaussian_ORIGINAL, dim=1) + 1e-5) | |
| ORIGINAL_sentence_term_loss = self.term_bce_loss(term_out_ORIGINAL.squeeze(1), sentence_level_target_term) | |
| ORIGINAL_sentence_loc_loss = torch.mean(loss_gaussian_ORIGINAL) | |
| ORIGINAL_sentence_touch_loss = self.mdn_bce_loss(eos_ORIGINAL.squeeze(1), sentence_level_target_eos) | |
| ORIGINAL_sentence_termination_loss.append(ORIGINAL_sentence_term_loss) | |
| ORIGINAL_sentence_loc_reconstruct_loss.append(ORIGINAL_sentence_loc_loss) | |
| ORIGINAL_sentence_touch_reconstruct_loss.append(ORIGINAL_sentence_touch_loss) | |
| # TYPE A | |
| if self.TYPE_A: | |
| sentence_C1 = sentence_Cs_1[sentence_batch_id] | |
| # sentence_Wc_rec_TYPE_A = torch.bmm(sentence_Cs[sentence_batch_id], sentence_W_mean.repeat(sentence_Cs[sentence_batch_id].size(0),1).unsqueeze(2)).squeeze(-1) | |
| sentence_Wc_rec_TYPE_A = torch.bmm(sentence_C1, \ | |
| sentence_W_mean.repeat(sentence_C1.size(0),1).unsqueeze(2)).squeeze(-1) | |
| sentence_Wcs_reconstruct_TYPE_A.append(sentence_Wc_rec_TYPE_A) | |
| sentence_W_lstm_in_TYPE_A = [] | |
| curr_id = 0 | |
| for i in range(user_sentence_level_stroke_length[sentence_batch_id][0]): | |
| sentence_W_lstm_in_TYPE_A.append(sentence_Wc_rec_TYPE_A[curr_id]) | |
| if i in sentence_SPLITS[sentence_batch_id]: | |
| curr_id += 1 | |
| sentence_Wc_t_rec_TYPE_A = torch.stack(sentence_W_lstm_in_TYPE_A) | |
| sentence_gen_lstm2_in_TYPE_A = torch.cat([sentence_level_gen_encoded, sentence_Wc_t_rec_TYPE_A], -1) | |
| sentence_gen_lstm2_in_TYPE_A = sentence_gen_lstm2_in_TYPE_A.unsqueeze(0) | |
| sentence_gen_out_TYPE_A, (c,h) = self.gen_state_lstm2(sentence_gen_lstm2_in_TYPE_A) | |
| sentence_gen_out_TYPE_A = sentence_gen_out_TYPE_A.squeeze(0) | |
| mdn_out_TYPE_A = self.gen_state_fc2(sentence_gen_out_TYPE_A) | |
| eos_TYPE_A = mdn_out_TYPE_A[:,0:1] | |
| [mu1_TYPE_A, mu2_TYPE_A, sig1_TYPE_A, sig2_TYPE_A, rho_TYPE_A, pi_TYPE_A] = torch.split(mdn_out_TYPE_A[:,1:], self.num_mixtures, 1) | |
| sig1_TYPE_A = sig1_TYPE_A.exp() + 1e-3 | |
| sig2_TYPE_A = sig2_TYPE_A.exp() + 1e-3 | |
| rho_TYPE_A = self.mdn_tanh(rho_TYPE_A) | |
| pi_TYPE_A = self.mdn_softmax(pi_TYPE_A) | |
| term_out_TYPE_A = self.term_fc1(sentence_gen_out_TYPE_A) | |
| term_out_TYPE_A = self.term_relu1(term_out_TYPE_A) | |
| term_out_TYPE_A = self.term_fc2(term_out_TYPE_A) | |
| term_out_TYPE_A = self.term_relu2(term_out_TYPE_A) | |
| term_out_TYPE_A = self.term_fc3(term_out_TYPE_A) | |
| term_pred_TYPE_A = self.term_sigmoid(term_out_TYPE_A) | |
| gaussian_TYPE_A = gaussian_2d(sentence_level_target_x, sentence_level_target_y, mu1_TYPE_A, mu2_TYPE_A, sig1_TYPE_A, sig2_TYPE_A, rho_TYPE_A) | |
| loss_gaussian_TYPE_A = - torch.log(torch.sum(pi_TYPE_A*gaussian_TYPE_A, dim=1) + 1e-5) | |
| TYPE_A_sentence_term_loss = self.term_bce_loss(term_out_TYPE_A.squeeze(1), sentence_level_target_term) | |
| TYPE_A_sentence_loc_loss = torch.mean(loss_gaussian_TYPE_A) | |
| TYPE_A_sentence_touch_loss = self.mdn_bce_loss(eos_TYPE_A.squeeze(1), sentence_level_target_eos) | |
| TYPE_A_sentence_termination_loss.append(TYPE_A_sentence_term_loss) | |
| TYPE_A_sentence_loc_reconstruct_loss.append(TYPE_A_sentence_loc_loss) | |
| TYPE_A_sentence_touch_reconstruct_loss.append(TYPE_A_sentence_touch_loss) | |
| # TYPE B | |
| if self.TYPE_B: | |
| unique_char_matrix_1 = sentence_unique_char_matrices_1[sentence_batch_id] | |
| unique_char_matrices_1 = torch.stack(unique_char_matrix_1) | |
| unique_char_matrices_1 = unique_char_matrices_1.squeeze(1) | |
| # sentence_W_c_TYPE_B_RAW = torch.bmm(unique_char_matrices, sentence_W_mean.repeat(unique_char_matrices.size(0), 1).unsqueeze(2)).squeeze(-1) | |
| sentence_W_c_TYPE_B_RAW = torch.bmm(unique_char_matrices_1, | |
| sentence_W_mean.repeat(unique_char_matrices_1.size(0), 1).unsqueeze(2)).squeeze(-1) | |
| sentence_W_c_TYPE_B_RAW = sentence_W_c_TYPE_B_RAW.unsqueeze(0) | |
| sentence_Wc_rec_TYPE_B, (c,h) = self.magic_lstm(sentence_W_c_TYPE_B_RAW) | |
| sentence_Wc_rec_TYPE_B = sentence_Wc_rec_TYPE_B.squeeze(0) | |
| sentence_Wcs_reconstruct_TYPE_B.append(sentence_Wc_rec_TYPE_B) | |
| sentence_W_lstm_in_TYPE_B = [] | |
| curr_id = 0 | |
| for i in range(user_sentence_level_stroke_length[sentence_batch_id][0]): | |
| sentence_W_lstm_in_TYPE_B.append(sentence_Wc_rec_TYPE_B[curr_id]) | |
| if i in sentence_SPLITS[sentence_batch_id]: | |
| curr_id += 1 | |
| sentence_Wc_t_rec_TYPE_B = torch.stack(sentence_W_lstm_in_TYPE_B) | |
| sentence_gen_lstm2_in_TYPE_B = torch.cat([sentence_level_gen_encoded, sentence_Wc_t_rec_TYPE_B], -1) | |
| sentence_gen_lstm2_in_TYPE_B = sentence_gen_lstm2_in_TYPE_B.unsqueeze(0) | |
| sentence_gen_out_TYPE_B, (c,h) = self.gen_state_lstm2(sentence_gen_lstm2_in_TYPE_B) | |
| sentence_gen_out_TYPE_B = sentence_gen_out_TYPE_B.squeeze(0) | |
| mdn_out_TYPE_B = self.gen_state_fc2(sentence_gen_out_TYPE_B) | |
| eos_TYPE_B = mdn_out_TYPE_B[:,0:1] | |
| [mu1_TYPE_B, mu2_TYPE_B, sig1_TYPE_B, sig2_TYPE_B, rho_TYPE_B, pi_TYPE_B] = torch.split(mdn_out_TYPE_B[:,1:], self.num_mixtures, 1) | |
| sig1_TYPE_B = sig1_TYPE_B.exp() + 1e-3 | |
| sig2_TYPE_B = sig2_TYPE_B.exp() + 1e-3 | |
| rho_TYPE_B = self.mdn_tanh(rho_TYPE_B) | |
| pi_TYPE_B = self.mdn_softmax(pi_TYPE_B) | |
| term_out_TYPE_B = self.term_fc1(sentence_gen_out_TYPE_B) | |
| term_out_TYPE_B = self.term_relu1(term_out_TYPE_B) | |
| term_out_TYPE_B = self.term_fc2(term_out_TYPE_B) | |
| term_out_TYPE_B = self.term_relu2(term_out_TYPE_B) | |
| term_out_TYPE_B = self.term_fc3(term_out_TYPE_B) | |
| term_pred_TYPE_B = self.term_sigmoid(term_out_TYPE_B) | |
| gaussian_TYPE_B = gaussian_2d(sentence_level_target_x, sentence_level_target_y, mu1_TYPE_B, mu2_TYPE_B, sig1_TYPE_B, sig2_TYPE_B, rho_TYPE_B) | |
| loss_gaussian_TYPE_B = - torch.log(torch.sum(pi_TYPE_B*gaussian_TYPE_B, dim=1) + 1e-5) | |
| TYPE_B_sentence_term_loss = self.term_bce_loss(term_out_TYPE_B.squeeze(1), sentence_level_target_term) | |
| TYPE_B_sentence_loc_loss = torch.mean(loss_gaussian_TYPE_B) | |
| TYPE_B_sentence_touch_loss = self.mdn_bce_loss(eos_TYPE_B.squeeze(1), sentence_level_target_eos) | |
| TYPE_B_sentence_termination_loss.append(TYPE_B_sentence_term_loss) | |
| TYPE_B_sentence_loc_reconstruct_loss.append(TYPE_B_sentence_loc_loss) | |
| TYPE_B_sentence_touch_reconstruct_loss.append(TYPE_B_sentence_touch_loss) | |
| if self.ORIGINAL: | |
| ALL_ORIGINAL_sentence_termination_loss.append(torch.mean(torch.stack(ORIGINAL_sentence_termination_loss))) | |
| ALL_ORIGINAL_sentence_loc_reconstruct_loss.append(torch.mean(torch.stack(ORIGINAL_sentence_loc_reconstruct_loss))) | |
| ALL_ORIGINAL_sentence_touch_reconstruct_loss.append(torch.mean(torch.stack(ORIGINAL_sentence_touch_reconstruct_loss))) | |
| if self.TYPE_A: | |
| ALL_TYPE_A_sentence_termination_loss.append(torch.mean(torch.stack(TYPE_A_sentence_termination_loss))) | |
| ALL_TYPE_A_sentence_loc_reconstruct_loss.append(torch.mean(torch.stack(TYPE_A_sentence_loc_reconstruct_loss))) | |
| ALL_TYPE_A_sentence_touch_reconstruct_loss.append(torch.mean(torch.stack(TYPE_A_sentence_touch_reconstruct_loss))) | |
| if self.REC: | |
| TYPE_A_sentence_WC_reconstruct_loss = [] | |
| for sentence_batch_id in range(len(sentence_Wc_rec_TYPE_)): | |
| sentence_Wc_ORIGINAL = sentence_Wc_rec_TYPE_[sentence_batch_id] | |
| sentence_Wc_TYPE_A = sentence_Wcs_reconstruct_TYPE_A[sentence_batch_id] | |
| sentence_WC_reconstruct_loss_TYPE_A = torch.mean(torch.mean(torch.mul(sentence_Wc_ORIGINAL - sentence_Wc_TYPE_A, sentence_Wc_ORIGINAL - sentence_Wc_TYPE_A), -1)) | |
| TYPE_A_sentence_WC_reconstruct_loss.append(sentence_WC_reconstruct_loss_TYPE_A) | |
| ALL_TYPE_A_sentence_WC_reconstruct_loss.append(torch.mean(torch.stack(TYPE_A_sentence_WC_reconstruct_loss))) | |
| if self.TYPE_B: | |
| ALL_TYPE_B_sentence_termination_loss.append(torch.mean(torch.stack(TYPE_B_sentence_termination_loss))) | |
| ALL_TYPE_B_sentence_loc_reconstruct_loss.append(torch.mean(torch.stack(TYPE_B_sentence_loc_reconstruct_loss))) | |
| ALL_TYPE_B_sentence_touch_reconstruct_loss.append(torch.mean(torch.stack(TYPE_B_sentence_touch_reconstruct_loss))) | |
| if self.REC: | |
| TYPE_B_sentence_WC_reconstruct_loss = [] | |
| for sentence_batch_id in range(len(sentence_Wc_rec_TYPE_)): | |
| sentence_Wc_ORIGINAL = sentence_Wc_rec_TYPE_[sentence_batch_id] | |
| sentence_Wc_TYPE_B = sentence_Wcs_reconstruct_TYPE_B[sentence_batch_id] | |
| sentence_WC_reconstruct_loss_TYPE_B = torch.mean(torch.mean(torch.mul(sentence_Wc_ORIGINAL - sentence_Wc_TYPE_B, sentence_Wc_ORIGINAL - sentence_Wc_TYPE_B), -1)) | |
| TYPE_B_sentence_WC_reconstruct_loss.append(sentence_WC_reconstruct_loss_TYPE_B) | |
| ALL_TYPE_B_sentence_WC_reconstruct_loss.append(torch.mean(torch.stack(TYPE_B_sentence_WC_reconstruct_loss))) | |
| if self.word_loss: | |
| user_word_level_stroke_in = word_level_stroke_in[uid] | |
| user_word_level_stroke_out = word_level_stroke_out[uid] | |
| user_word_level_stroke_length = word_level_stroke_length[uid] | |
| user_word_level_term = word_level_term[uid] | |
| user_word_level_char = word_level_char[uid] | |
| user_word_level_char_length = word_level_char_length[uid] | |
| word_batch_size = len(user_word_level_stroke_in) | |
| word_inf_state_out = self.inf_state_fc1(user_word_level_stroke_out) | |
| word_inf_state_out = self.inf_state_relu(word_inf_state_out) | |
| word_inf_state_out, (c,h) = self.inf_state_lstm(word_inf_state_out) | |
| word_gen_state_out = self.gen_state_fc1(user_word_level_stroke_in) | |
| word_gen_state_out = self.gen_state_relu(word_gen_state_out) | |
| word_gen_state_out, (c,h) = self.gen_state_lstm1(word_gen_state_out) | |
| word_Ws = [] | |
| word_Wc_rec_ORIGINAL = [] | |
| word_SPLITS = [] | |
| word_Cs_1 = [] | |
| word_unique_char_matrices_1 = [] | |
| W_C_ORIGINALS = [] | |
| for word_batch_id in range(word_batch_size): | |
| curr_seq_len = user_word_level_stroke_length[word_batch_id][0] | |
| curr_char_len = user_word_level_char_length[word_batch_id][0] | |
| char_vector = torch.eye(len(CHARACTERS))[user_word_level_char[word_batch_id][:curr_char_len]].to(self.device) | |
| current_term = user_word_level_term[word_batch_id][:curr_seq_len].unsqueeze(-1) | |
| split_ids = torch.nonzero(current_term)[:,0] | |
| char_vector_1 = self.char_vec_fc_1(char_vector) | |
| char_vector_1 = self.char_vec_relu_1(char_vector_1) | |
| unique_char_matrices_1 = [] | |
| for cid in range(len(char_vector)): | |
| # Tower 1 | |
| unique_char_vector_1 = char_vector_1[cid:cid+1] | |
| unique_char_input_1 = unique_char_vector_1.unsqueeze(0) | |
| unique_char_out_1, (c,h) = self.char_lstm_1(unique_char_input_1) | |
| unique_char_out_1 = unique_char_out_1.squeeze(0) | |
| unique_char_out_1 = self.char_vec_fc2_1(unique_char_out_1) | |
| unique_char_matrix_1 = unique_char_out_1.view([-1,1,self.weight_dim,self.weight_dim]) | |
| unique_char_matrix_1 = unique_char_matrix_1.squeeze(1) | |
| unique_char_matrices_1.append(unique_char_matrix_1) | |
| # Tower 1 | |
| char_out_1 = char_vector_1.unsqueeze(0) | |
| char_out_1, (c,h) = self.char_lstm_1(char_out_1) | |
| char_out_1 = char_out_1.squeeze(0) | |
| char_out_1 = self.char_vec_fc2_1(char_out_1) | |
| char_matrix_1 = char_out_1.view([-1,1,self.weight_dim,self.weight_dim]) | |
| char_matrix_1 = char_matrix_1.squeeze(1) | |
| char_matrix_inv_1 = torch.inverse(char_matrix_1) | |
| W_c_t = word_inf_state_out[word_batch_id][:curr_seq_len] | |
| W_c = torch.stack([W_c_t[i] for i in split_ids]) | |
| W_C_ORIGINAL = {} | |
| for i in range(curr_char_len): | |
| sub_s = "".join(CHARACTERS[i] for i in user_word_level_char[word_batch_id][:i+1]) | |
| W_C_ORIGINAL[sub_s] = [W_c[i]] | |
| W_C_ORIGINALS.append(W_C_ORIGINAL) | |
| # W = torch.bmm(char_matrix_inv, W_c.unsqueeze(2)).squeeze(-1) | |
| W = torch.bmm(char_matrix_inv_1, | |
| W_c.unsqueeze(2)).squeeze(-1) | |
| word_Ws.append(W) | |
| word_Wc_rec_ORIGINAL.append(W_c) | |
| word_SPLITS.append(split_ids) | |
| # word_Cs.append(char_matrix) | |
| # word_unique_char_matrices.append(unique_char_matrices) | |
| word_Cs_1.append(char_matrix_1) | |
| word_unique_char_matrices_1.append(unique_char_matrices_1) | |
| word_Ws_stacked = torch.cat(word_Ws, 0) | |
| word_Ws_reshaped = word_Ws_stacked.view([-1,self.weight_dim]) | |
| word_W_mean = word_Ws_reshaped.mean(0) | |
| word_Ws_reshaped_mean_repeat = word_W_mean.repeat(word_Ws_reshaped.size(0),1) | |
| word_Ws_consistency_loss = torch.mean(torch.mean(torch.mul(word_Ws_reshaped_mean_repeat - word_Ws_reshaped, word_Ws_reshaped_mean_repeat - word_Ws_reshaped), -1)) | |
| ALL_word_W_consistency_loss.append(word_Ws_consistency_loss) | |
| # word | |
| ORIGINAL_word_termination_loss = [] | |
| ORIGINAL_word_loc_reconstruct_loss = [] | |
| ORIGINAL_word_touch_reconstruct_loss = [] | |
| TYPE_A_word_termination_loss = [] | |
| TYPE_A_word_loc_reconstruct_loss = [] | |
| TYPE_A_word_touch_reconstruct_loss = [] | |
| TYPE_B_word_termination_loss = [] | |
| TYPE_B_word_loc_reconstruct_loss = [] | |
| TYPE_B_word_touch_reconstruct_loss = [] | |
| TYPE_C_word_termination_loss = [] | |
| TYPE_C_word_loc_reconstruct_loss = [] | |
| TYPE_C_word_touch_reconstruct_loss = [] | |
| TYPE_D_word_termination_loss = [] | |
| TYPE_D_word_loc_reconstruct_loss = [] | |
| TYPE_D_word_touch_reconstruct_loss = [] | |
| word_Wcs_reconstruct_TYPE_A = [] | |
| word_Wcs_reconstruct_TYPE_B = [] | |
| word_Wcs_reconstruct_TYPE_C = [] | |
| word_Wcs_reconstruct_TYPE_D = [] | |
| # segment | |
| ALL_segment_W_consistency_loss = [] | |
| ALL_ORIGINAL_segment_termination_loss = [] | |
| ALL_ORIGINAL_segment_loc_reconstruct_loss = [] | |
| ALL_ORIGINAL_segment_touch_reconstruct_loss = [] | |
| ALL_TYPE_A_segment_termination_loss = [] | |
| ALL_TYPE_A_segment_loc_reconstruct_loss = [] | |
| ALL_TYPE_A_segment_touch_reconstruct_loss = [] | |
| ALL_TYPE_A_segment_WC_reconstruct_loss = [] | |
| ALL_TYPE_B_segment_termination_loss = [] | |
| ALL_TYPE_B_segment_loc_reconstruct_loss = [] | |
| ALL_TYPE_B_segment_touch_reconstruct_loss = [] | |
| ALL_TYPE_B_segment_WC_reconstruct_loss = [] | |
| ALL_segment_Wcs_reconstruct_TYPE_A = [] | |
| ALL_segment_Wcs_reconstruct_TYPE_B = [] | |
| W_C_SEGMENTS = [] | |
| W_C_UNIQUES = [] | |
| for word_batch_id in range(word_batch_size): | |
| word_level_gen_encoded = word_gen_state_out[word_batch_id][:user_word_level_stroke_length[word_batch_id][0]] | |
| word_level_target_eos = user_word_level_stroke_out[word_batch_id][:user_word_level_stroke_length[word_batch_id][0]][:,2] | |
| word_level_target_x = user_word_level_stroke_out[word_batch_id][:user_word_level_stroke_length[word_batch_id][0]][:,0:1] | |
| word_level_target_y = user_word_level_stroke_out[word_batch_id][:user_word_level_stroke_length[word_batch_id][0]][:,1:2] | |
| word_level_target_term = user_word_level_term[word_batch_id][:user_word_level_stroke_length[word_batch_id][0]] | |
| # ORIGINAL | |
| if self.ORIGINAL: | |
| word_W_lstm_in_ORIGINAL = [] | |
| curr_id = 0 | |
| for i in range(user_word_level_stroke_length[word_batch_id][0]): | |
| word_W_lstm_in_ORIGINAL.append(word_Wc_rec_ORIGINAL[word_batch_id][curr_id]) | |
| if i in word_SPLITS[word_batch_id]: | |
| curr_id += 1 | |
| word_W_lstm_in_ORIGINAL = torch.stack(word_W_lstm_in_ORIGINAL) | |
| word_Wc_t_ORIGINAL = word_W_lstm_in_ORIGINAL | |
| word_gen_lstm2_in_ORIGINAL = torch.cat([word_level_gen_encoded, word_Wc_t_ORIGINAL], -1) | |
| word_gen_lstm2_in_ORIGINAL = word_gen_lstm2_in_ORIGINAL.unsqueeze(0) | |
| word_gen_out_ORIGINAL,(c,h) = self.gen_state_lstm2(word_gen_lstm2_in_ORIGINAL) | |
| word_gen_out_ORIGINAL = word_gen_out_ORIGINAL.squeeze(0) | |
| mdn_out_ORIGINAL = self.gen_state_fc2(word_gen_out_ORIGINAL) | |
| eos_ORIGINAL = mdn_out_ORIGINAL[:,0:1] | |
| [mu1_ORIGINAL, mu2_ORIGINAL, sig1_ORIGINAL, sig2_ORIGINAL, rho_ORIGINAL, pi_ORIGINAL] = torch.split(mdn_out_ORIGINAL[:,1:], self.num_mixtures, 1) | |
| sig1_ORIGINAL = sig1_ORIGINAL.exp() + 1e-3 | |
| sig2_ORIGINAL = sig2_ORIGINAL.exp() + 1e-3 | |
| rho_ORIGINAL = self.mdn_tanh(rho_ORIGINAL) | |
| pi_ORIGINAL = self.mdn_softmax(pi_ORIGINAL) | |
| term_out_ORIGINAL = self.term_fc1(word_gen_out_ORIGINAL) | |
| term_out_ORIGINAL = self.term_relu1(term_out_ORIGINAL) | |
| term_out_ORIGINAL = self.term_fc2(term_out_ORIGINAL) | |
| term_out_ORIGINAL = self.term_relu2(term_out_ORIGINAL) | |
| term_out_ORIGINAL = self.term_fc3(term_out_ORIGINAL) | |
| term_pred_ORIGINAL = self.term_sigmoid(term_out_ORIGINAL) | |
| gaussian_ORIGINAL = gaussian_2d(word_level_target_x, word_level_target_y, mu1_ORIGINAL, mu2_ORIGINAL, sig1_ORIGINAL, sig2_ORIGINAL, rho_ORIGINAL) | |
| loss_gaussian_ORIGINAL = - torch.log(torch.sum(pi_ORIGINAL*gaussian_ORIGINAL, dim=1) + 1e-5) | |
| ORIGINAL_word_term_loss = self.term_bce_loss(term_out_ORIGINAL.squeeze(1), word_level_target_term) | |
| ORIGINAL_word_loc_loss = torch.mean(loss_gaussian_ORIGINAL) | |
| ORIGINAL_word_touch_loss = self.mdn_bce_loss(eos_ORIGINAL.squeeze(1), word_level_target_eos) | |
| ORIGINAL_word_termination_loss.append(ORIGINAL_word_term_loss) | |
| ORIGINAL_word_loc_reconstruct_loss.append(ORIGINAL_word_loc_loss) | |
| ORIGINAL_word_touch_reconstruct_loss.append(ORIGINAL_word_touch_loss) | |
| # TYPE A | |
| if self.TYPE_A: | |
| word_C1 = word_Cs_1[word_batch_id] | |
| word_Wc_rec_TYPE_A = torch.bmm(word_C1, | |
| word_W_mean.repeat(word_C1.size(0),1).unsqueeze(2)).squeeze(-1) | |
| word_Wcs_reconstruct_TYPE_A.append(word_Wc_rec_TYPE_A) | |
| word_W_lstm_in_TYPE_A = [] | |
| curr_id = 0 | |
| for i in range(user_word_level_stroke_length[word_batch_id][0]): | |
| word_W_lstm_in_TYPE_A.append(word_Wc_rec_TYPE_A[curr_id]) | |
| if i in word_SPLITS[word_batch_id]: | |
| curr_id += 1 | |
| word_Wc_t_rec_TYPE_A = torch.stack(word_W_lstm_in_TYPE_A) | |
| word_gen_lstm2_in_TYPE_A = torch.cat([word_level_gen_encoded, word_Wc_t_rec_TYPE_A], -1) | |
| word_gen_lstm2_in_TYPE_A = word_gen_lstm2_in_TYPE_A.unsqueeze(0) | |
| word_gen_out_TYPE_A, (c,h) = self.gen_state_lstm2(word_gen_lstm2_in_TYPE_A) | |
| word_gen_out_TYPE_A = word_gen_out_TYPE_A.squeeze(0) | |
| mdn_out_TYPE_A = self.gen_state_fc2(word_gen_out_TYPE_A) | |
| eos_TYPE_A = mdn_out_TYPE_A[:,0:1] | |
| [mu1_TYPE_A, mu2_TYPE_A, sig1_TYPE_A, sig2_TYPE_A, rho_TYPE_A, pi_TYPE_A] = torch.split(mdn_out_TYPE_A[:,1:], self.num_mixtures, 1) | |
| sig1_TYPE_A = sig1_TYPE_A.exp() + 1e-3 | |
| sig2_TYPE_A = sig2_TYPE_A.exp() + 1e-3 | |
| rho_TYPE_A = self.mdn_tanh(rho_TYPE_A) | |
| pi_TYPE_A = self.mdn_softmax(pi_TYPE_A) | |
| term_out_TYPE_A = self.term_fc1(word_gen_out_TYPE_A) | |
| term_out_TYPE_A = self.term_relu1(term_out_TYPE_A) | |
| term_out_TYPE_A = self.term_fc2(term_out_TYPE_A) | |
| term_out_TYPE_A = self.term_relu2(term_out_TYPE_A) | |
| term_out_TYPE_A = self.term_fc3(term_out_TYPE_A) | |
| term_pred_TYPE_A = self.term_sigmoid(term_out_TYPE_A) | |
| gaussian_TYPE_A = gaussian_2d(word_level_target_x, word_level_target_y, mu1_TYPE_A, mu2_TYPE_A, sig1_TYPE_A, sig2_TYPE_A, rho_TYPE_A) | |
| loss_gaussian_TYPE_A = - torch.log(torch.sum(pi_TYPE_A*gaussian_TYPE_A, dim=1) + 1e-5) | |
| TYPE_A_word_term_loss = self.term_bce_loss(term_out_TYPE_A.squeeze(1), word_level_target_term) | |
| TYPE_A_word_loc_loss = torch.mean(loss_gaussian_TYPE_A) | |
| TYPE_A_word_touch_loss = self.mdn_bce_loss(eos_TYPE_A.squeeze(1), word_level_target_eos) | |
| TYPE_A_word_termination_loss.append(TYPE_A_word_term_loss) | |
| TYPE_A_word_loc_reconstruct_loss.append(TYPE_A_word_loc_loss) | |
| TYPE_A_word_touch_reconstruct_loss.append(TYPE_A_word_touch_loss) | |
| # TYPE B | |
| if self.TYPE_B: | |
| unique_char_matrix_1 = word_unique_char_matrices_1[word_batch_id] | |
| unique_char_matrices_1 = torch.stack(unique_char_matrix_1) | |
| unique_char_matrices_1 = unique_char_matrices_1.squeeze(1) | |
| # word_W_c_TYPE_B_RAW = torch.bmm(unique_char_matrices, word_W_mean.repeat(unique_char_matrices.size(0), 1).unsqueeze(2)).squeeze(-1) | |
| word_W_c_TYPE_B_RAW = torch.bmm(unique_char_matrices_1, | |
| word_W_mean.repeat(unique_char_matrices_1.size(0), 1).unsqueeze(2)).squeeze(-1) | |
| word_W_c_TYPE_B_RAW = word_W_c_TYPE_B_RAW.unsqueeze(0) | |
| word_Wc_rec_TYPE_B, (c,h) = self.magic_lstm(word_W_c_TYPE_B_RAW) | |
| word_Wc_rec_TYPE_B = word_Wc_rec_TYPE_B.squeeze(0) | |
| word_Wcs_reconstruct_TYPE_B.append(word_Wc_rec_TYPE_B) | |
| word_W_lstm_in_TYPE_B = [] | |
| curr_id = 0 | |
| for i in range(user_word_level_stroke_length[word_batch_id][0]): | |
| word_W_lstm_in_TYPE_B.append(word_Wc_rec_TYPE_B[curr_id]) | |
| if i in word_SPLITS[word_batch_id]: | |
| curr_id += 1 | |
| word_Wc_t_rec_TYPE_B = torch.stack(word_W_lstm_in_TYPE_B) | |
| word_gen_lstm2_in_TYPE_B = torch.cat([word_level_gen_encoded, word_Wc_t_rec_TYPE_B], -1) | |
| word_gen_lstm2_in_TYPE_B = word_gen_lstm2_in_TYPE_B.unsqueeze(0) | |
| word_gen_out_TYPE_B, (c,h) = self.gen_state_lstm2(word_gen_lstm2_in_TYPE_B) | |
| word_gen_out_TYPE_B = word_gen_out_TYPE_B.squeeze(0) | |
| mdn_out_TYPE_B = self.gen_state_fc2(word_gen_out_TYPE_B) | |
| eos_TYPE_B = mdn_out_TYPE_B[:,0:1] | |
| [mu1_TYPE_B, mu2_TYPE_B, sig1_TYPE_B, sig2_TYPE_B, rho_TYPE_B, pi_TYPE_B] = torch.split(mdn_out_TYPE_B[:,1:], self.num_mixtures, 1) | |
| sig1_TYPE_B = sig1_TYPE_B.exp() + 1e-3 | |
| sig2_TYPE_B = sig2_TYPE_B.exp() + 1e-3 | |
| rho_TYPE_B = self.mdn_tanh(rho_TYPE_B) | |
| pi_TYPE_B = self.mdn_softmax(pi_TYPE_B) | |
| term_out_TYPE_B = self.term_fc1(word_gen_out_TYPE_B) | |
| term_out_TYPE_B = self.term_relu1(term_out_TYPE_B) | |
| term_out_TYPE_B = self.term_fc2(term_out_TYPE_B) | |
| term_out_TYPE_B = self.term_relu2(term_out_TYPE_B) | |
| term_out_TYPE_B = self.term_fc3(term_out_TYPE_B) | |
| term_pred_TYPE_B = self.term_sigmoid(term_out_TYPE_B) | |
| gaussian_TYPE_B = gaussian_2d(word_level_target_x, word_level_target_y, mu1_TYPE_B, mu2_TYPE_B, sig1_TYPE_B, sig2_TYPE_B, rho_TYPE_B) | |
| loss_gaussian_TYPE_B = - torch.log(torch.sum(pi_TYPE_B*gaussian_TYPE_B, dim=1) + 1e-5) | |
| TYPE_B_word_term_loss = self.term_bce_loss(term_out_TYPE_B.squeeze(1), word_level_target_term) | |
| TYPE_B_word_loc_loss = torch.mean(loss_gaussian_TYPE_B) | |
| TYPE_B_word_touch_loss = self.mdn_bce_loss(eos_TYPE_B.squeeze(1), word_level_target_eos) | |
| TYPE_B_word_termination_loss.append(TYPE_B_word_term_loss) | |
| TYPE_B_word_loc_reconstruct_loss.append(TYPE_B_word_loc_loss) | |
| TYPE_B_word_touch_reconstruct_loss.append(TYPE_B_word_touch_loss) | |
| # TYPE C | |
| # if self.TYPE_C: | |
| user_segment_level_stroke_in = segment_level_stroke_in[uid][word_batch_id] | |
| user_segment_level_stroke_out = segment_level_stroke_out[uid][word_batch_id] | |
| user_segment_level_stroke_length = segment_level_stroke_length[uid][word_batch_id] | |
| user_segment_level_term = segment_level_term[uid][word_batch_id] | |
| user_segment_level_char = segment_level_char[uid][word_batch_id] | |
| user_segment_level_char_length = segment_level_char_length[uid][word_batch_id] | |
| segment_batch_size = len(user_segment_level_stroke_in) | |
| segment_inf_state_out = self.inf_state_fc1(user_segment_level_stroke_out) | |
| segment_inf_state_out = self.inf_state_relu(segment_inf_state_out) | |
| segment_inf_state_out, (c,h) = self.inf_state_lstm(segment_inf_state_out) | |
| segment_gen_state_out = self.gen_state_fc1(user_segment_level_stroke_in) | |
| segment_gen_state_out = self.gen_state_relu(segment_gen_state_out) | |
| segment_gen_state_out, (c,h) = self.gen_state_lstm1(segment_gen_state_out) | |
| segment_Ws = [] | |
| segment_Wc_rec_ORIGINAL = [] | |
| segment_SPLITS = [] | |
| segment_Cs_1 = [] | |
| segment_unique_char_matrices_1 = [] | |
| W_C_SEGMENT = {} | |
| for segment_batch_id in range(segment_batch_size): | |
| curr_seq_len = user_segment_level_stroke_length[segment_batch_id][0] | |
| curr_char_len = user_segment_level_char_length[segment_batch_id][0] | |
| char_vector = torch.eye(len(CHARACTERS))[user_segment_level_char[segment_batch_id][:curr_char_len]].to(self.device) | |
| current_term = user_segment_level_term[segment_batch_id][:curr_seq_len].unsqueeze(-1) | |
| split_ids = torch.nonzero(current_term)[:,0] | |
| char_vector_1 = self.char_vec_fc_1(char_vector) | |
| char_vector_1 = self.char_vec_relu_1(char_vector_1) | |
| unique_char_matrices_1 = [] | |
| for cid in range(len(char_vector)): | |
| # Tower 1 | |
| unique_char_vector_1 = char_vector_1[cid:cid+1] | |
| unique_char_input_1 = unique_char_vector_1.unsqueeze(0) | |
| unique_char_out_1, (c,h) = self.char_lstm_1(unique_char_input_1) | |
| unique_char_out_1 = unique_char_out_1.squeeze(0) | |
| unique_char_out_1 = self.char_vec_fc2_1(unique_char_out_1) | |
| unique_char_matrix_1 = unique_char_out_1.view([-1,1,self.weight_dim,self.weight_dim]) | |
| unique_char_matrix_1 = unique_char_matrix_1.squeeze(1) | |
| unique_char_matrices_1.append(unique_char_matrix_1) | |
| # Tower 1 | |
| char_out_1 = char_vector_1.unsqueeze(0) | |
| char_out_1, (c,h) = self.char_lstm_1(char_out_1) | |
| char_out_1 = char_out_1.squeeze(0) | |
| char_out_1 = self.char_vec_fc2_1(char_out_1) | |
| char_matrix_1 = char_out_1.view([-1,1,self.weight_dim,self.weight_dim]) | |
| char_matrix_1 = char_matrix_1.squeeze(1) | |
| char_matrix_inv_1 = torch.inverse(char_matrix_1) | |
| W_c_t = segment_inf_state_out[segment_batch_id][:curr_seq_len] | |
| W_c = torch.stack([W_c_t[i] for i in split_ids]) | |
| for i in range(curr_char_len): | |
| sub_s = "".join(CHARACTERS[i] for i in user_segment_level_char[segment_batch_id][:i+1]) | |
| if sub_s in W_C_SEGMENT: | |
| W_C_SEGMENT[sub_s].append(W_c[i]) | |
| else: | |
| W_C_SEGMENT[sub_s] = [W_c[i]] | |
| W = torch.bmm(char_matrix_inv_1, | |
| W_c.unsqueeze(2)).squeeze(-1) | |
| segment_Ws.append(W) | |
| segment_Wc_rec_ORIGINAL.append(W_c) | |
| segment_SPLITS.append(split_ids) | |
| segment_Cs_1.append(char_matrix_1) | |
| segment_unique_char_matrices_1.append(unique_char_matrices_1) | |
| W_C_SEGMENTS.append(W_C_SEGMENT) | |
| if self.segment_loss: | |
| segment_Ws_stacked = torch.cat(segment_Ws, 0) | |
| segment_Ws_reshaped = segment_Ws_stacked.view([-1,self.weight_dim]) | |
| segment_W_mean = segment_Ws_reshaped.mean(0) | |
| segment_Ws_reshaped_mean_repeat = segment_W_mean.repeat(segment_Ws_reshaped.size(0),1) | |
| segment_Ws_consistency_loss = torch.mean(torch.mean(torch.mul(segment_Ws_reshaped_mean_repeat - segment_Ws_reshaped, segment_Ws_reshaped_mean_repeat - segment_Ws_reshaped), -1)) | |
| ALL_segment_W_consistency_loss.append(segment_Ws_consistency_loss) | |
| ORIGINAL_segment_termination_loss = [] | |
| ORIGINAL_segment_loc_reconstruct_loss = [] | |
| ORIGINAL_segment_touch_reconstruct_loss = [] | |
| TYPE_A_segment_termination_loss = [] | |
| TYPE_A_segment_loc_reconstruct_loss = [] | |
| TYPE_A_segment_touch_reconstruct_loss = [] | |
| TYPE_B_segment_termination_loss = [] | |
| TYPE_B_segment_loc_reconstruct_loss = [] | |
| TYPE_B_segment_touch_reconstruct_loss = [] | |
| segment_Wcs_reconstruct_TYPE_A = [] | |
| segment_Wcs_reconstruct_TYPE_B = [] | |
| for segment_batch_id in range(segment_batch_size): | |
| segment_level_gen_encoded = segment_gen_state_out[segment_batch_id][:user_segment_level_stroke_length[segment_batch_id][0]] | |
| segment_level_target_eos = user_segment_level_stroke_out[segment_batch_id][:user_segment_level_stroke_length[segment_batch_id][0]][:,2] | |
| segment_level_target_x = user_segment_level_stroke_out[segment_batch_id][:user_segment_level_stroke_length[segment_batch_id][0]][:,0:1] | |
| segment_level_target_y = user_segment_level_stroke_out[segment_batch_id][:user_segment_level_stroke_length[segment_batch_id][0]][:,1:2] | |
| segment_level_target_term = user_segment_level_term[segment_batch_id][:user_segment_level_stroke_length[segment_batch_id][0]] | |
| if self.ORIGINAL: | |
| segment_W_lstm_in_ORIGINAL = [] | |
| curr_id = 0 | |
| for i in range(user_segment_level_stroke_length[segment_batch_id][0]): | |
| segment_W_lstm_in_ORIGINAL.append(segment_Wc_rec_ORIGINAL[segment_batch_id][curr_id]) | |
| if i in segment_SPLITS[segment_batch_id]: | |
| curr_id += 1 | |
| segment_W_lstm_in_ORIGINAL = torch.stack(segment_W_lstm_in_ORIGINAL) | |
| segment_Wc_t_ORIGINAL = segment_W_lstm_in_ORIGINAL | |
| segment_gen_lstm2_in_ORIGINAL = torch.cat([segment_level_gen_encoded, segment_Wc_t_ORIGINAL], -1) | |
| segment_gen_lstm2_in_ORIGINAL = segment_gen_lstm2_in_ORIGINAL.unsqueeze(0) | |
| segment_gen_out_ORIGINAL,(c,h) = self.gen_state_lstm2(segment_gen_lstm2_in_ORIGINAL) | |
| segment_gen_out_ORIGINAL = segment_gen_out_ORIGINAL.squeeze(0) | |
| mdn_out_ORIGINAL = self.gen_state_fc2(segment_gen_out_ORIGINAL) | |
| eos_ORIGINAL = mdn_out_ORIGINAL[:,0:1] | |
| [mu1_ORIGINAL, mu2_ORIGINAL, sig1_ORIGINAL, sig2_ORIGINAL, rho_ORIGINAL, pi_ORIGINAL] = torch.split(mdn_out_ORIGINAL[:,1:], self.num_mixtures, 1) | |
| sig1_ORIGINAL = sig1_ORIGINAL.exp() + 1e-3 | |
| sig2_ORIGINAL = sig2_ORIGINAL.exp() + 1e-3 | |
| rho_ORIGINAL = self.mdn_tanh(rho_ORIGINAL) | |
| pi_ORIGINAL = self.mdn_softmax(pi_ORIGINAL) | |
| term_out_ORIGINAL = self.term_fc1(segment_gen_out_ORIGINAL) | |
| term_out_ORIGINAL = self.term_relu1(term_out_ORIGINAL) | |
| term_out_ORIGINAL = self.term_fc2(term_out_ORIGINAL) | |
| term_out_ORIGINAL = self.term_relu2(term_out_ORIGINAL) | |
| term_out_ORIGINAL = self.term_fc3(term_out_ORIGINAL) | |
| term_pred_ORIGINAL = self.term_sigmoid(term_out_ORIGINAL) | |
| gaussian_ORIGINAL = gaussian_2d(segment_level_target_x, segment_level_target_y, mu1_ORIGINAL, mu2_ORIGINAL, sig1_ORIGINAL, sig2_ORIGINAL, rho_ORIGINAL) | |
| loss_gaussian_ORIGINAL = - torch.log(torch.sum(pi_ORIGINAL*gaussian_ORIGINAL, dim=1) + 1e-5) | |
| ORIGINAL_segment_term_loss = self.term_bce_loss(term_out_ORIGINAL.squeeze(1), segment_level_target_term) | |
| ORIGINAL_segment_loc_loss = torch.mean(loss_gaussian_ORIGINAL) | |
| ORIGINAL_segment_touch_loss = self.mdn_bce_loss(eos_ORIGINAL.squeeze(1), segment_level_target_eos) | |
| ORIGINAL_segment_termination_loss.append(ORIGINAL_segment_term_loss) | |
| ORIGINAL_segment_loc_reconstruct_loss.append(ORIGINAL_segment_loc_loss) | |
| ORIGINAL_segment_touch_reconstruct_loss.append(ORIGINAL_segment_touch_loss) | |
| # TYPE A | |
| if self.TYPE_A: | |
| segment_C1 = segment_Cs_1[segment_batch_id] | |
| segment_Wc_rec_TYPE_A = torch.bmm(segment_C1, | |
| segment_W_mean.repeat(segment_C1.size(0),1).unsqueeze(2)).squeeze(-1) | |
| segment_Wcs_reconstruct_TYPE_A.append(segment_Wc_rec_TYPE_A) | |
| segment_W_lstm_in_TYPE_A = [] | |
| curr_id = 0 | |
| for i in range(user_segment_level_stroke_length[segment_batch_id][0]): | |
| segment_W_lstm_in_TYPE_A.append(segment_Wc_rec_TYPE_A[curr_id]) | |
| if i in segment_SPLITS[segment_batch_id]: | |
| curr_id += 1 | |
| segment_Wc_t_rec_TYPE_A = torch.stack(segment_W_lstm_in_TYPE_A) | |
| segment_gen_lstm2_in_TYPE_A = torch.cat([segment_level_gen_encoded, segment_Wc_t_rec_TYPE_A], -1) | |
| segment_gen_lstm2_in_TYPE_A = segment_gen_lstm2_in_TYPE_A.unsqueeze(0) | |
| segment_gen_out_TYPE_A, (c,h) = self.gen_state_lstm2(segment_gen_lstm2_in_TYPE_A) | |
| segment_gen_out_TYPE_A = segment_gen_out_TYPE_A.squeeze(0) | |
| mdn_out_TYPE_A = self.gen_state_fc2(segment_gen_out_TYPE_A) | |
| eos_TYPE_A = mdn_out_TYPE_A[:,0:1] | |
| [mu1_TYPE_A, mu2_TYPE_A, sig1_TYPE_A, sig2_TYPE_A, rho_TYPE_A, pi_TYPE_A] = torch.split(mdn_out_TYPE_A[:,1:], self.num_mixtures, 1) | |
| sig1_TYPE_A = sig1_TYPE_A.exp() + 1e-3 | |
| sig2_TYPE_A = sig2_TYPE_A.exp() + 1e-3 | |
| rho_TYPE_A = self.mdn_tanh(rho_TYPE_A) | |
| pi_TYPE_A = self.mdn_softmax(pi_TYPE_A) | |
| term_out_TYPE_A = self.term_fc1(segment_gen_out_TYPE_A) | |
| term_out_TYPE_A = self.term_relu1(term_out_TYPE_A) | |
| term_out_TYPE_A = self.term_fc2(term_out_TYPE_A) | |
| term_out_TYPE_A = self.term_relu2(term_out_TYPE_A) | |
| term_out_TYPE_A = self.term_fc3(term_out_TYPE_A) | |
| term_pred_TYPE_A = self.term_sigmoid(term_out_TYPE_A) | |
| gaussian_TYPE_A = gaussian_2d(segment_level_target_x, segment_level_target_y, mu1_TYPE_A, mu2_TYPE_A, sig1_TYPE_A, sig2_TYPE_A, rho_TYPE_A) | |
| loss_gaussian_TYPE_A = - torch.log(torch.sum(pi_TYPE_A*gaussian_TYPE_A, dim=1) + 1e-5) | |
| TYPE_A_segment_term_loss = self.term_bce_loss(term_out_TYPE_A.squeeze(1), segment_level_target_term) | |
| TYPE_A_segment_loc_loss = torch.mean(loss_gaussian_TYPE_A) | |
| TYPE_A_segment_touch_loss = self.mdn_bce_loss(eos_TYPE_A.squeeze(1), segment_level_target_eos) | |
| TYPE_A_segment_termination_loss.append(TYPE_A_segment_term_loss) | |
| TYPE_A_segment_loc_reconstruct_loss.append(TYPE_A_segment_loc_loss) | |
| TYPE_A_segment_touch_reconstruct_loss.append(TYPE_A_segment_touch_loss) | |
| # TYPE B | |
| if self.TYPE_B: | |
| unique_char_matrix_1 = segment_unique_char_matrices_1[segment_batch_id] | |
| unique_char_matrices_1 = torch.stack(unique_char_matrix_1) | |
| unique_char_matrices_1 = unique_char_matrices_1.squeeze(1) | |
| # segment_W_c_TYPE_B_RAW = torch.bmm(unique_char_matrices, segment_W_mean.repeat(unique_char_matrices.size(0), 1).unsqueeze(2)).squeeze(-1) | |
| segment_W_c_TYPE_B_RAW = torch.bmm(unique_char_matrices_1, | |
| segment_W_mean.repeat(unique_char_matrices_1.size(0), 1).unsqueeze(2)).squeeze(-1) | |
| segment_W_c_TYPE_B_RAW = segment_W_c_TYPE_B_RAW.unsqueeze(0) | |
| segment_Wc_rec_TYPE_B, (c,h) = self.magic_lstm(segment_W_c_TYPE_B_RAW) | |
| segment_Wc_rec_TYPE_B = segment_Wc_rec_TYPE_B.squeeze(0) | |
| segment_Wcs_reconstruct_TYPE_B.append(segment_Wc_rec_TYPE_B) | |
| segment_W_lstm_in_TYPE_B = [] | |
| curr_id = 0 | |
| for i in range(user_segment_level_stroke_length[segment_batch_id][0]): | |
| segment_W_lstm_in_TYPE_B.append(segment_Wc_rec_TYPE_B[curr_id]) | |
| if i in segment_SPLITS[segment_batch_id]: | |
| curr_id += 1 | |
| segment_Wc_t_rec_TYPE_B = torch.stack(segment_W_lstm_in_TYPE_B) | |
| segment_gen_lstm2_in_TYPE_B = torch.cat([segment_level_gen_encoded, segment_Wc_t_rec_TYPE_B], -1) | |
| segment_gen_lstm2_in_TYPE_B = segment_gen_lstm2_in_TYPE_B.unsqueeze(0) | |
| segment_gen_out_TYPE_B, (c,h) = self.gen_state_lstm2(segment_gen_lstm2_in_TYPE_B) | |
| segment_gen_out_TYPE_B = segment_gen_out_TYPE_B.squeeze(0) | |
| mdn_out_TYPE_B = self.gen_state_fc2(segment_gen_out_TYPE_B) | |
| eos_TYPE_B = mdn_out_TYPE_B[:,0:1] | |
| [mu1_TYPE_B, mu2_TYPE_B, sig1_TYPE_B, sig2_TYPE_B, rho_TYPE_B, pi_TYPE_B] = torch.split(mdn_out_TYPE_B[:,1:], self.num_mixtures, 1) | |
| sig1_TYPE_B = sig1_TYPE_B.exp() + 1e-3 | |
| sig2_TYPE_B = sig2_TYPE_B.exp() + 1e-3 | |
| rho_TYPE_B = self.mdn_tanh(rho_TYPE_B) | |
| pi_TYPE_B = self.mdn_softmax(pi_TYPE_B) | |
| term_out_TYPE_B = self.term_fc1(segment_gen_out_TYPE_B) | |
| term_out_TYPE_B = self.term_relu1(term_out_TYPE_B) | |
| term_out_TYPE_B = self.term_fc2(term_out_TYPE_B) | |
| term_out_TYPE_B = self.term_relu2(term_out_TYPE_B) | |
| term_out_TYPE_B = self.term_fc3(term_out_TYPE_B) | |
| term_pred_TYPE_B = self.term_sigmoid(term_out_TYPE_B) | |
| gaussian_TYPE_B = gaussian_2d(segment_level_target_x, segment_level_target_y, mu1_TYPE_B, mu2_TYPE_B, sig1_TYPE_B, sig2_TYPE_B, rho_TYPE_B) | |
| loss_gaussian_TYPE_B = - torch.log(torch.sum(pi_TYPE_B*gaussian_TYPE_B, dim=1) + 1e-5) | |
| TYPE_B_segment_term_loss = self.term_bce_loss(term_out_TYPE_B.squeeze(1), segment_level_target_term) | |
| TYPE_B_segment_loc_loss = torch.mean(loss_gaussian_TYPE_B) | |
| TYPE_B_segment_touch_loss = self.mdn_bce_loss(eos_TYPE_B.squeeze(1), segment_level_target_eos) | |
| TYPE_B_segment_termination_loss.append(TYPE_B_segment_term_loss) | |
| TYPE_B_segment_loc_reconstruct_loss.append(TYPE_B_segment_loc_loss) | |
| TYPE_B_segment_touch_reconstruct_loss.append(TYPE_B_segment_touch_loss) | |
| if self.ORIGINAL: | |
| ALL_ORIGINAL_segment_termination_loss.append(torch.mean(torch.stack(ORIGINAL_segment_termination_loss))) | |
| ALL_ORIGINAL_segment_loc_reconstruct_loss.append(torch.mean(torch.stack(ORIGINAL_segment_loc_reconstruct_loss))) | |
| ALL_ORIGINAL_segment_touch_reconstruct_loss.append(torch.mean(torch.stack(ORIGINAL_segment_touch_reconstruct_loss))) | |
| if self.TYPE_A: | |
| ALL_TYPE_A_segment_termination_loss.append(torch.mean(torch.stack(TYPE_A_segment_termination_loss))) | |
| ALL_TYPE_A_segment_loc_reconstruct_loss.append(torch.mean(torch.stack(TYPE_A_segment_loc_reconstruct_loss))) | |
| ALL_TYPE_A_segment_touch_reconstruct_loss.append(torch.mean(torch.stack(TYPE_A_segment_touch_reconstruct_loss))) | |
| if self.REC: | |
| TYPE_A_segment_WC_reconstruct_loss = [] | |
| for segment_batch_id in range(len(segment_Wc_rec_ORIGINAL)): | |
| segment_Wc_ORIGINAL = segment_Wc_rec_ORIGINAL[segment_batch_id] | |
| segment_Wc_TYPE_A = segment_Wcs_reconstruct_TYPE_A[segment_batch_id] | |
| segment_WC_reconstruct_loss_TYPE_A = torch.mean(torch.mean(torch.mul(segment_Wc_ORIGINAL - segment_Wc_TYPE_A, segment_Wc_ORIGINAL - segment_Wc_TYPE_A), -1)) | |
| TYPE_A_segment_WC_reconstruct_loss.append(segment_WC_reconstruct_loss_TYPE_A) | |
| ALL_TYPE_A_segment_WC_reconstruct_loss.append(torch.mean(torch.stack(TYPE_A_segment_WC_reconstruct_loss))) | |
| if self.TYPE_B: | |
| ALL_TYPE_B_segment_termination_loss.append(torch.mean(torch.stack(TYPE_B_segment_termination_loss))) | |
| ALL_TYPE_B_segment_loc_reconstruct_loss.append(torch.mean(torch.stack(TYPE_B_segment_loc_reconstruct_loss))) | |
| ALL_TYPE_B_segment_touch_reconstruct_loss.append(torch.mean(torch.stack(TYPE_B_segment_touch_reconstruct_loss))) | |
| if self.REC: | |
| TYPE_B_segment_WC_reconstruct_loss = [] | |
| for segment_batch_id in range(len(segment_Wc_rec_ORIGINAL)): | |
| segment_Wc_ORIGINAL = segment_Wc_rec_ORIGINAL[segment_batch_id] | |
| segment_Wc_TYPE_B = segment_Wcs_reconstruct_TYPE_B[segment_batch_id] | |
| segment_WC_reconstruct_loss_TYPE_B = torch.mean(torch.mean(torch.mul(segment_Wc_ORIGINAL - segment_Wc_TYPE_B, segment_Wc_ORIGINAL - segment_Wc_TYPE_B), -1)) | |
| TYPE_B_segment_WC_reconstruct_loss.append(segment_WC_reconstruct_loss_TYPE_B) | |
| ALL_TYPE_B_segment_WC_reconstruct_loss.append(torch.mean(torch.stack(TYPE_B_segment_WC_reconstruct_loss))) | |
| if self.TYPE_C: | |
| # target | |
| original_W_c = word_Wc_rec_ORIGINAL[word_batch_id] | |
| word_Wc_rec_TYPE_C = [] | |
| for segment_batch_id in range(len(segment_Wc_rec_ORIGINAL)): | |
| if segment_batch_id == 0: | |
| for each_segment_Wc in segment_Wc_rec_ORIGINAL[segment_batch_id]: | |
| word_Wc_rec_TYPE_C.append(each_segment_Wc) | |
| prev_id = len(word_Wc_rec_TYPE_C) - 1 | |
| else: | |
| prev_original_W_c = original_W_c[prev_id] | |
| for each_segment_Wc in segment_Wc_rec_ORIGINAL[segment_batch_id]: | |
| magic_inp = torch.stack([prev_original_W_c, each_segment_Wc]) | |
| magic_inp = magic_inp.unsqueeze(0) | |
| type_c_out, (c,h) = self.magic_lstm(magic_inp) | |
| type_c_out = type_c_out.squeeze(0) | |
| word_Wc_rec_TYPE_C.append(type_c_out[-1]) | |
| prev_id = len(word_Wc_rec_TYPE_C) - 1 | |
| word_Wc_rec_TYPE_C = torch.stack(word_Wc_rec_TYPE_C) | |
| word_Wcs_reconstruct_TYPE_C.append(word_Wc_rec_TYPE_C) | |
| if len(word_Wc_rec_TYPE_C) == len(word_SPLITS[word_batch_id]): | |
| word_W_lstm_in_TYPE_C = [] | |
| curr_id = 0 | |
| for i in range(user_word_level_stroke_length[word_batch_id][0]): | |
| word_W_lstm_in_TYPE_C.append(word_Wc_rec_TYPE_C[curr_id]) | |
| if i in word_SPLITS[word_batch_id]: | |
| curr_id += 1 | |
| word_Wc_t_rec_TYPE_C = torch.stack(word_W_lstm_in_TYPE_C) | |
| word_gen_lstm2_in_TYPE_C = torch.cat([word_level_gen_encoded, word_Wc_t_rec_TYPE_C], -1) | |
| word_gen_lstm2_in_TYPE_C = word_gen_lstm2_in_TYPE_C.unsqueeze(0) | |
| word_gen_out_TYPE_C, (c,h) = self.gen_state_lstm2(word_gen_lstm2_in_TYPE_C) | |
| word_gen_out_TYPE_C = word_gen_out_TYPE_C.squeeze(0) | |
| mdn_out_TYPE_C = self.gen_state_fc2(word_gen_out_TYPE_C) | |
| eos_TYPE_C = mdn_out_TYPE_C[:,0:1] | |
| [mu1_TYPE_C, mu2_TYPE_C, sig1_TYPE_C, sig2_TYPE_C, rho_TYPE_C, pi_TYPE_C] = torch.split(mdn_out_TYPE_C[:,1:], self.num_mixtures, 1) | |
| sig1_TYPE_C = sig1_TYPE_C.exp() + 1e-3 | |
| sig2_TYPE_C = sig2_TYPE_C.exp() + 1e-3 | |
| rho_TYPE_C = self.mdn_tanh(rho_TYPE_C) | |
| pi_TYPE_C = self.mdn_softmax(pi_TYPE_C) | |
| term_out_TYPE_C = self.term_fc1(word_gen_out_TYPE_C) | |
| term_out_TYPE_C = self.term_relu1(term_out_TYPE_C) | |
| term_out_TYPE_C = self.term_fc2(term_out_TYPE_C) | |
| term_out_TYPE_C = self.term_relu2(term_out_TYPE_C) | |
| term_out_TYPE_C = self.term_fc3(term_out_TYPE_C) | |
| term_pred_TYPE_C = self.term_sigmoid(term_out_TYPE_C) | |
| gaussian_TYPE_C = gaussian_2d(word_level_target_x, word_level_target_y, mu1_TYPE_C, mu2_TYPE_C, sig1_TYPE_C, sig2_TYPE_C, rho_TYPE_C) | |
| loss_gaussian_TYPE_C = - torch.log(torch.sum(pi_TYPE_C*gaussian_TYPE_C, dim=1) + 1e-5) | |
| TYPE_C_word_term_loss = self.term_bce_loss(term_out_TYPE_C.squeeze(1), word_level_target_term) | |
| TYPE_C_word_loc_loss = torch.mean(loss_gaussian_TYPE_C) | |
| TYPE_C_word_touch_loss = self.mdn_bce_loss(eos_TYPE_C.squeeze(1), word_level_target_eos) | |
| TYPE_C_word_termination_loss.append(TYPE_C_word_term_loss) | |
| TYPE_C_word_loc_reconstruct_loss.append(TYPE_C_word_loc_loss) | |
| TYPE_C_word_touch_reconstruct_loss.append(TYPE_C_word_touch_loss) | |
| else: | |
| print ("not C") | |
| if self.TYPE_D: | |
| word_Wc_rec_TYPE_D = [] | |
| TYPE_D_REF = [] | |
| for segment_batch_id in range(len(segment_Wc_rec_ORIGINAL)): | |
| if segment_batch_id == 0: | |
| for each_segment_Wc in segment_Wc_rec_ORIGINAL[segment_batch_id]: | |
| word_Wc_rec_TYPE_D.append(each_segment_Wc) | |
| TYPE_D_REF.append(segment_Wc_rec_ORIGINAL[segment_batch_id][-1]) | |
| else: | |
| for each_segment_Wc in segment_Wc_rec_ORIGINAL[segment_batch_id]: | |
| magic_inp = torch.cat([torch.stack(TYPE_D_REF, 0), each_segment_Wc.unsqueeze(0)], 0) | |
| magic_inp = magic_inp.unsqueeze(0) | |
| TYPE_D_out, (c,h) = self.magic_lstm(magic_inp) | |
| TYPE_D_out = TYPE_D_out.squeeze(0) | |
| word_Wc_rec_TYPE_D.append(TYPE_D_out[-1]) | |
| TYPE_D_REF.append(segment_Wc_rec_ORIGINAL[segment_batch_id][-1]) | |
| word_Wc_rec_TYPE_D = torch.stack(word_Wc_rec_TYPE_D) | |
| word_Wcs_reconstruct_TYPE_D.append(word_Wc_rec_TYPE_D) | |
| if len(word_Wc_rec_TYPE_D) == len(word_SPLITS[word_batch_id]): | |
| word_W_lstm_in_TYPE_D = [] | |
| curr_id = 0 | |
| for i in range(user_word_level_stroke_length[word_batch_id][0]): | |
| word_W_lstm_in_TYPE_D.append(word_Wc_rec_TYPE_D[curr_id]) | |
| if i in word_SPLITS[word_batch_id]: | |
| curr_id += 1 | |
| word_Wc_t_rec_TYPE_D = torch.stack(word_W_lstm_in_TYPE_D) | |
| word_gen_lstm2_in_TYPE_D = torch.cat([word_level_gen_encoded, word_Wc_t_rec_TYPE_D], -1) | |
| word_gen_lstm2_in_TYPE_D = word_gen_lstm2_in_TYPE_D.unsqueeze(0) | |
| word_gen_out_TYPE_D, (c,h) = self.gen_state_lstm2(word_gen_lstm2_in_TYPE_D) | |
| word_gen_out_TYPE_D = word_gen_out_TYPE_D.squeeze(0) | |
| mdn_out_TYPE_D = self.gen_state_fc2(word_gen_out_TYPE_D) | |
| eos_TYPE_D = mdn_out_TYPE_D[:,0:1] | |
| [mu1_TYPE_D, mu2_TYPE_D, sig1_TYPE_D, sig2_TYPE_D, rho_TYPE_D, pi_TYPE_D] = torch.split(mdn_out_TYPE_D[:,1:], self.num_mixtures, 1) | |
| sig1_TYPE_D = sig1_TYPE_D.exp() + 1e-3 | |
| sig2_TYPE_D = sig2_TYPE_D.exp() + 1e-3 | |
| rho_TYPE_D = self.mdn_tanh(rho_TYPE_D) | |
| pi_TYPE_D = self.mdn_softmax(pi_TYPE_D) | |
| term_out_TYPE_D = self.term_fc1(word_gen_out_TYPE_D) | |
| term_out_TYPE_D = self.term_relu1(term_out_TYPE_D) | |
| term_out_TYPE_D = self.term_fc2(term_out_TYPE_D) | |
| term_out_TYPE_D = self.term_relu2(term_out_TYPE_D) | |
| term_out_TYPE_D = self.term_fc3(term_out_TYPE_D) | |
| term_pred_TYPE_D = self.term_sigmoid(term_out_TYPE_D) | |
| gaussian_TYPE_D = gaussian_2d(word_level_target_x, word_level_target_y, mu1_TYPE_D, mu2_TYPE_D, sig1_TYPE_D, sig2_TYPE_D, rho_TYPE_D) | |
| loss_gaussian_TYPE_D = - torch.log(torch.sum(pi_TYPE_D*gaussian_TYPE_D, dim=1) + 1e-5) | |
| TYPE_D_word_term_loss = self.term_bce_loss(term_out_TYPE_D.squeeze(1), word_level_target_term) | |
| TYPE_D_word_loc_loss = torch.mean(loss_gaussian_TYPE_D) | |
| TYPE_D_word_touch_loss = self.mdn_bce_loss(eos_TYPE_D.squeeze(1), word_level_target_eos) | |
| TYPE_D_word_termination_loss.append(TYPE_D_word_term_loss) | |
| TYPE_D_word_loc_reconstruct_loss.append(TYPE_D_word_loc_loss) | |
| TYPE_D_word_touch_reconstruct_loss.append(TYPE_D_word_touch_loss) | |
| else: | |
| print ("not D") | |
| # word | |
| if self.ORIGINAL: | |
| ALL_ORIGINAL_word_termination_loss.append(torch.mean(torch.stack(ORIGINAL_word_termination_loss))) | |
| ALL_ORIGINAL_word_loc_reconstruct_loss.append(torch.mean(torch.stack(ORIGINAL_word_loc_reconstruct_loss))) | |
| ALL_ORIGINAL_word_touch_reconstruct_loss.append(torch.mean(torch.stack(ORIGINAL_word_touch_reconstruct_loss))) | |
| if self.TYPE_A: | |
| ALL_TYPE_A_word_termination_loss.append(torch.mean(torch.stack(TYPE_A_word_termination_loss))) | |
| ALL_TYPE_A_word_loc_reconstruct_loss.append(torch.mean(torch.stack(TYPE_A_word_loc_reconstruct_loss))) | |
| ALL_TYPE_A_word_touch_reconstruct_loss.append(torch.mean(torch.stack(TYPE_A_word_touch_reconstruct_loss))) | |
| if self.REC: | |
| TYPE_A_word_WC_reconstruct_loss = [] | |
| for word_batch_id in range(len(word_Wc_rec_ORIGINAL)): | |
| word_Wc_ORIGINAL = word_Wc_rec_ORIGINAL[word_batch_id] | |
| word_Wc_TYPE_A = word_Wcs_reconstruct_TYPE_A[word_batch_id] | |
| if len(word_Wc_ORIGINAL) == len(word_Wc_TYPE_A): | |
| word_WC_reconstruct_loss_TYPE_A = torch.mean(torch.mean(torch.mul(word_Wc_ORIGINAL - word_Wc_TYPE_A, word_Wc_ORIGINAL - word_Wc_TYPE_A), -1)) | |
| TYPE_A_word_WC_reconstruct_loss.append(word_WC_reconstruct_loss_TYPE_A) | |
| if len(TYPE_A_word_WC_reconstruct_loss) > 0: | |
| ALL_TYPE_A_word_WC_reconstruct_loss.append(torch.mean(torch.stack(TYPE_A_word_WC_reconstruct_loss))) | |
| if self.TYPE_B: | |
| ALL_TYPE_B_word_termination_loss.append(torch.mean(torch.stack(TYPE_B_word_termination_loss))) | |
| ALL_TYPE_B_word_loc_reconstruct_loss.append(torch.mean(torch.stack(TYPE_B_word_loc_reconstruct_loss))) | |
| ALL_TYPE_B_word_touch_reconstruct_loss.append(torch.mean(torch.stack(TYPE_B_word_touch_reconstruct_loss))) | |
| if self.REC: | |
| TYPE_B_word_WC_reconstruct_loss = [] | |
| for word_batch_id in range(len(word_Wc_rec_ORIGINAL)): | |
| word_Wc_ORIGINAL = word_Wc_rec_ORIGINAL[word_batch_id] | |
| word_Wc_TYPE_B = word_Wcs_reconstruct_TYPE_B[word_batch_id] | |
| if len(word_Wc_ORIGINAL) == len(word_Wc_TYPE_B): | |
| word_WC_reconstruct_loss_TYPE_B = torch.mean(torch.mean(torch.mul(word_Wc_ORIGINAL - word_Wc_TYPE_B, word_Wc_ORIGINAL - word_Wc_TYPE_B), -1)) | |
| TYPE_B_word_WC_reconstruct_loss.append(word_WC_reconstruct_loss_TYPE_B) | |
| if len(TYPE_B_word_WC_reconstruct_loss) > 0: | |
| ALL_TYPE_B_word_WC_reconstruct_loss.append(torch.mean(torch.stack(TYPE_B_word_WC_reconstruct_loss))) | |
| if self.TYPE_C: | |
| ALL_TYPE_C_word_termination_loss.append(torch.mean(torch.stack(TYPE_C_word_termination_loss))) | |
| ALL_TYPE_C_word_loc_reconstruct_loss.append(torch.mean(torch.stack(TYPE_C_word_loc_reconstruct_loss))) | |
| ALL_TYPE_C_word_touch_reconstruct_loss.append(torch.mean(torch.stack(TYPE_C_word_touch_reconstruct_loss))) | |
| if self.REC: | |
| TYPE_C_word_WC_reconstruct_loss = [] | |
| for word_batch_id in range(len(word_Wc_rec_ORIGINAL)): | |
| word_Wc_ORIGINAL = word_Wc_rec_ORIGINAL[word_batch_id] | |
| word_Wc_TYPE_C = word_Wcs_reconstruct_TYPE_C[word_batch_id] | |
| if len(word_Wc_ORIGINAL) == len(word_Wc_TYPE_C): | |
| word_WC_reconstruct_loss_TYPE_C = torch.mean(torch.mean(torch.mul(word_Wc_ORIGINAL - word_Wc_TYPE_C, word_Wc_ORIGINAL - word_Wc_TYPE_C), -1)) | |
| TYPE_C_word_WC_reconstruct_loss.append(word_WC_reconstruct_loss_TYPE_C) | |
| if len(TYPE_C_word_WC_reconstruct_loss) > 0: | |
| ALL_TYPE_C_word_WC_reconstruct_loss.append(torch.mean(torch.stack(TYPE_C_word_WC_reconstruct_loss))) | |
| if self.TYPE_D: | |
| ALL_TYPE_D_word_termination_loss.append(torch.mean(torch.stack(TYPE_D_word_termination_loss))) | |
| ALL_TYPE_D_word_loc_reconstruct_loss.append(torch.mean(torch.stack(TYPE_D_word_loc_reconstruct_loss))) | |
| ALL_TYPE_D_word_touch_reconstruct_loss.append(torch.mean(torch.stack(TYPE_D_word_touch_reconstruct_loss))) | |
| if self.REC: | |
| TYPE_D_word_WC_reconstruct_loss = [] | |
| for word_batch_id in range(len(word_Wc_rec_ORIGINAL)): | |
| word_Wc_ORIGINAL = word_Wc_rec_ORIGINAL[word_batch_id] | |
| word_Wc_TYPE_D = word_Wcs_reconstruct_TYPE_D[word_batch_id] | |
| if len(word_Wc_ORIGINAL) == len(word_Wc_TYPE_D): | |
| word_WC_reconstruct_loss_TYPE_D = torch.mean(torch.mean(torch.mul(word_Wc_ORIGINAL - word_Wc_TYPE_D, word_Wc_ORIGINAL - word_Wc_TYPE_D), -1)) | |
| TYPE_D_word_WC_reconstruct_loss.append(word_WC_reconstruct_loss_TYPE_D) | |
| if len(TYPE_D_word_WC_reconstruct_loss) > 0: | |
| ALL_TYPE_D_word_WC_reconstruct_loss.append(torch.mean(torch.stack(TYPE_D_word_WC_reconstruct_loss))) | |
| # segment | |
| if self.segment_loss: | |
| SUPER_ALL_segment_W_consistency_loss.append(torch.mean(torch.stack(ALL_segment_W_consistency_loss))) | |
| if self.ORIGINAL: | |
| SUPER_ALL_ORIGINAL_segment_termination_loss.append(torch.mean(torch.stack(ALL_ORIGINAL_segment_termination_loss))) | |
| SUPER_ALL_ORIGINAL_segment_loc_reconstruct_loss.append(torch.mean(torch.stack(ALL_ORIGINAL_segment_loc_reconstruct_loss))) | |
| SUPER_ALL_ORIGINAL_segment_touch_reconstruct_loss.append(torch.mean(torch.stack(ALL_ORIGINAL_segment_touch_reconstruct_loss))) | |
| if self.TYPE_A: | |
| SUPER_ALL_TYPE_A_segment_termination_loss.append(torch.mean(torch.stack(ALL_TYPE_A_segment_termination_loss))) | |
| SUPER_ALL_TYPE_A_segment_loc_reconstruct_loss.append(torch.mean(torch.stack(ALL_TYPE_A_segment_loc_reconstruct_loss))) | |
| SUPER_ALL_TYPE_A_segment_touch_reconstruct_loss.append(torch.mean(torch.stack(ALL_TYPE_A_segment_touch_reconstruct_loss))) | |
| if self.REC: | |
| SUPER_ALL_TYPE_A_segment_WC_reconstruct_loss.append(torch.mean(torch.stack(ALL_TYPE_A_segment_WC_reconstruct_loss))) | |
| if self.TYPE_B: | |
| SUPER_ALL_TYPE_B_segment_termination_loss.append(torch.mean(torch.stack(ALL_TYPE_B_segment_termination_loss))) | |
| SUPER_ALL_TYPE_B_segment_loc_reconstruct_loss.append(torch.mean(torch.stack(ALL_TYPE_B_segment_loc_reconstruct_loss))) | |
| SUPER_ALL_TYPE_B_segment_touch_reconstruct_loss.append(torch.mean(torch.stack(ALL_TYPE_B_segment_touch_reconstruct_loss))) | |
| if self.REC: | |
| SUPER_ALL_TYPE_B_segment_WC_reconstruct_loss.append(torch.mean(torch.stack(ALL_TYPE_B_segment_WC_reconstruct_loss))) | |
| total_sentence_loss = 0 | |
| sentence_losses = [] | |
| if self.sentence_loss: | |
| mean_ORIGINAL_sentence_termination_loss = 0 | |
| mean_ORIGINAL_sentence_loc_reconstruct_loss = 0 | |
| mean_ORIGINAL_sentence_touch_reconstruct_loss = 0 | |
| mean_TYPE_A_sentence_termination_loss = 0 | |
| mean_TYPE_A_sentence_loc_reconstruct_loss = 0 | |
| mean_TYPE_A_sentence_touch_reconstruct_loss = 0 | |
| mean_TYPE_B_sentence_termination_loss = 0 | |
| mean_TYPE_B_sentence_loc_reconstruct_loss = 0 | |
| mean_TYPE_B_sentence_touch_reconstruct_loss = 0 | |
| mean_TYPE_A_sentence_WC_reconstruct_loss = 0 | |
| mean_TYPE_B_sentence_WC_reconstruct_loss = 0 | |
| mean_sentence_W_consistency_loss = torch.mean(torch.stack(ALL_sentence_W_consistency_loss)) | |
| if self.ORIGINAL: | |
| mean_ORIGINAL_sentence_termination_loss = torch.mean(torch.stack(ALL_ORIGINAL_sentence_termination_loss)) | |
| mean_ORIGINAL_sentence_loc_reconstruct_loss = torch.mean(torch.stack(ALL_ORIGINAL_sentence_loc_reconstruct_loss)) | |
| mean_ORIGINAL_sentence_touch_reconstruct_loss = torch.mean(torch.stack(ALL_ORIGINAL_sentence_touch_reconstruct_loss)) | |
| if self.TYPE_A: | |
| mean_TYPE_A_sentence_termination_loss = torch.mean(torch.stack(ALL_TYPE_A_sentence_termination_loss)) | |
| mean_TYPE_A_sentence_loc_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_A_sentence_loc_reconstruct_loss)) | |
| mean_TYPE_A_sentence_touch_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_A_sentence_touch_reconstruct_loss)) | |
| if self.REC: | |
| mean_TYPE_A_sentence_WC_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_A_sentence_WC_reconstruct_loss)) | |
| if self.TYPE_B: | |
| mean_TYPE_B_sentence_termination_loss = torch.mean(torch.stack(ALL_TYPE_B_sentence_termination_loss)) | |
| mean_TYPE_B_sentence_loc_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_B_sentence_loc_reconstruct_loss)) | |
| mean_TYPE_B_sentence_touch_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_B_sentence_touch_reconstruct_loss)) | |
| if self.REC: | |
| mean_TYPE_B_sentence_WC_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_B_sentence_WC_reconstruct_loss)) | |
| total_sentence_loss = mean_sentence_W_consistency_loss + mean_ORIGINAL_sentence_termination_loss + mean_ORIGINAL_sentence_loc_reconstruct_loss + mean_ORIGINAL_sentence_touch_reconstruct_loss + mean_TYPE_A_sentence_termination_loss + mean_TYPE_A_sentence_loc_reconstruct_loss + mean_TYPE_A_sentence_touch_reconstruct_loss + mean_TYPE_B_sentence_termination_loss + mean_TYPE_B_sentence_loc_reconstruct_loss + mean_TYPE_B_sentence_touch_reconstruct_loss + mean_TYPE_A_sentence_WC_reconstruct_loss + mean_TYPE_B_sentence_WC_reconstruct_loss | |
| sentence_losses = [total_sentence_loss, mean_sentence_W_consistency_loss, mean_ORIGINAL_sentence_termination_loss, mean_ORIGINAL_sentence_loc_reconstruct_loss, mean_ORIGINAL_sentence_touch_reconstruct_loss, mean_TYPE_A_sentence_termination_loss, mean_TYPE_A_sentence_loc_reconstruct_loss, mean_TYPE_A_sentence_touch_reconstruct_loss, mean_TYPE_B_sentence_termination_loss, mean_TYPE_B_sentence_loc_reconstruct_loss, mean_TYPE_B_sentence_touch_reconstruct_loss, mean_TYPE_A_sentence_WC_reconstruct_loss, mean_TYPE_B_sentence_WC_reconstruct_loss] | |
| total_word_loss = 0 | |
| word_losses = [] | |
| if self.word_loss: | |
| mean_ORIGINAL_word_termination_loss = 0 | |
| mean_ORIGINAL_word_loc_reconstruct_loss = 0 | |
| mean_ORIGINAL_word_touch_reconstruct_loss = 0 | |
| mean_TYPE_A_word_termination_loss = 0 | |
| mean_TYPE_A_word_loc_reconstruct_loss = 0 | |
| mean_TYPE_A_word_touch_reconstruct_loss = 0 | |
| mean_TYPE_B_word_termination_loss = 0 | |
| mean_TYPE_B_word_loc_reconstruct_loss = 0 | |
| mean_TYPE_B_word_touch_reconstruct_loss = 0 | |
| mean_TYPE_C_word_termination_loss = 0 | |
| mean_TYPE_C_word_loc_reconstruct_loss = 0 | |
| mean_TYPE_C_word_touch_reconstruct_loss = 0 | |
| mean_TYPE_D_word_termination_loss = 0 | |
| mean_TYPE_D_word_loc_reconstruct_loss = 0 | |
| mean_TYPE_D_word_touch_reconstruct_loss = 0 | |
| mean_TYPE_A_word_WC_reconstruct_loss = 0 | |
| mean_TYPE_B_word_WC_reconstruct_loss = 0 | |
| mean_TYPE_C_word_WC_reconstruct_loss = 0 | |
| mean_TYPE_D_word_WC_reconstruct_loss = 0 | |
| mean_word_W_consistency_loss = torch.mean(torch.stack(ALL_word_W_consistency_loss)) | |
| if self.ORIGINAL: | |
| mean_ORIGINAL_word_termination_loss = torch.mean(torch.stack(ALL_ORIGINAL_word_termination_loss)) | |
| mean_ORIGINAL_word_loc_reconstruct_loss = torch.mean(torch.stack(ALL_ORIGINAL_word_loc_reconstruct_loss)) | |
| mean_ORIGINAL_word_touch_reconstruct_loss = torch.mean(torch.stack(ALL_ORIGINAL_word_touch_reconstruct_loss)) | |
| if self.TYPE_A: | |
| mean_TYPE_A_word_termination_loss = torch.mean(torch.stack(ALL_TYPE_A_word_termination_loss)) | |
| mean_TYPE_A_word_loc_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_A_word_loc_reconstruct_loss)) | |
| mean_TYPE_A_word_touch_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_A_word_touch_reconstruct_loss)) | |
| if self.REC: | |
| mean_TYPE_A_word_WC_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_A_word_WC_reconstruct_loss)) | |
| if self.TYPE_B: | |
| mean_TYPE_B_word_termination_loss = torch.mean(torch.stack(ALL_TYPE_B_word_termination_loss)) | |
| mean_TYPE_B_word_loc_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_B_word_loc_reconstruct_loss)) | |
| mean_TYPE_B_word_touch_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_B_word_touch_reconstruct_loss)) | |
| if self.REC: | |
| mean_TYPE_B_word_WC_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_B_word_WC_reconstruct_loss)) | |
| if self.TYPE_C: | |
| mean_TYPE_C_word_termination_loss = torch.mean(torch.stack(ALL_TYPE_C_word_termination_loss)) | |
| mean_TYPE_C_word_loc_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_C_word_loc_reconstruct_loss)) | |
| mean_TYPE_C_word_touch_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_C_word_touch_reconstruct_loss)) | |
| if self.REC: | |
| mean_TYPE_C_word_WC_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_C_word_WC_reconstruct_loss)) | |
| if self.TYPE_D: | |
| mean_TYPE_D_word_termination_loss = torch.mean(torch.stack(ALL_TYPE_D_word_termination_loss)) | |
| mean_TYPE_D_word_loc_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_D_word_loc_reconstruct_loss)) | |
| mean_TYPE_D_word_touch_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_D_word_touch_reconstruct_loss)) | |
| if self.REC: | |
| mean_TYPE_D_word_WC_reconstruct_loss = torch.mean(torch.stack(ALL_TYPE_D_word_WC_reconstruct_loss)) | |
| total_word_loss = mean_word_W_consistency_loss + mean_ORIGINAL_word_termination_loss + mean_ORIGINAL_word_loc_reconstruct_loss + mean_ORIGINAL_word_touch_reconstruct_loss + mean_TYPE_A_word_termination_loss + mean_TYPE_A_word_loc_reconstruct_loss + mean_TYPE_A_word_touch_reconstruct_loss + mean_TYPE_B_word_termination_loss + mean_TYPE_B_word_loc_reconstruct_loss + mean_TYPE_B_word_touch_reconstruct_loss + mean_TYPE_C_word_termination_loss + mean_TYPE_C_word_loc_reconstruct_loss + mean_TYPE_C_word_touch_reconstruct_loss + mean_TYPE_D_word_termination_loss + mean_TYPE_D_word_loc_reconstruct_loss + mean_TYPE_D_word_touch_reconstruct_loss + mean_TYPE_A_word_WC_reconstruct_loss + mean_TYPE_B_word_WC_reconstruct_loss + mean_TYPE_C_word_WC_reconstruct_loss + mean_TYPE_D_word_WC_reconstruct_loss | |
| word_losses = [total_word_loss, mean_word_W_consistency_loss, mean_ORIGINAL_word_termination_loss, mean_ORIGINAL_word_loc_reconstruct_loss, mean_ORIGINAL_word_touch_reconstruct_loss, mean_TYPE_A_word_termination_loss, mean_TYPE_A_word_loc_reconstruct_loss, mean_TYPE_A_word_touch_reconstruct_loss, mean_TYPE_B_word_termination_loss, mean_TYPE_B_word_loc_reconstruct_loss, mean_TYPE_B_word_touch_reconstruct_loss, mean_TYPE_C_word_termination_loss, mean_TYPE_C_word_loc_reconstruct_loss, mean_TYPE_C_word_touch_reconstruct_loss, mean_TYPE_D_word_termination_loss, mean_TYPE_D_word_loc_reconstruct_loss, mean_TYPE_D_word_touch_reconstruct_loss, mean_TYPE_A_word_WC_reconstruct_loss, mean_TYPE_B_word_WC_reconstruct_loss, mean_TYPE_C_word_WC_reconstruct_loss, mean_TYPE_D_word_WC_reconstruct_loss] | |
| total_segment_loss = 0 | |
| segment_losses = [] | |
| if self.segment_loss: | |
| mean_segment_W_consistency_loss = torch.mean(torch.stack(SUPER_ALL_segment_W_consistency_loss)) | |
| mean_ORIGINAL_segment_termination_loss = 0 | |
| mean_ORIGINAL_segment_loc_reconstruct_loss = 0 | |
| mean_ORIGINAL_segment_touch_reconstruct_loss = 0 | |
| mean_TYPE_A_segment_termination_loss = 0 | |
| mean_TYPE_A_segment_loc_reconstruct_loss = 0 | |
| mean_TYPE_A_segment_touch_reconstruct_loss = 0 | |
| mean_TYPE_B_segment_termination_loss = 0 | |
| mean_TYPE_B_segment_loc_reconstruct_loss = 0 | |
| mean_TYPE_B_segment_touch_reconstruct_loss = 0 | |
| mean_TYPE_A_segment_WC_reconstruct_loss = 0 | |
| mean_TYPE_B_segment_WC_reconstruct_loss = 0 | |
| if self.ORIGINAL: | |
| mean_ORIGINAL_segment_termination_loss = torch.mean(torch.stack(SUPER_ALL_ORIGINAL_segment_termination_loss)) | |
| mean_ORIGINAL_segment_loc_reconstruct_loss = torch.mean(torch.stack(SUPER_ALL_ORIGINAL_segment_loc_reconstruct_loss)) | |
| mean_ORIGINAL_segment_touch_reconstruct_loss = torch.mean(torch.stack(SUPER_ALL_ORIGINAL_segment_touch_reconstruct_loss)) | |
| if self.TYPE_A: | |
| mean_TYPE_A_segment_termination_loss = torch.mean(torch.stack(SUPER_ALL_TYPE_A_segment_termination_loss)) | |
| mean_TYPE_A_segment_loc_reconstruct_loss = torch.mean(torch.stack(SUPER_ALL_TYPE_A_segment_loc_reconstruct_loss)) | |
| mean_TYPE_A_segment_touch_reconstruct_loss = torch.mean(torch.stack(SUPER_ALL_TYPE_A_segment_touch_reconstruct_loss)) | |
| if self.REC: | |
| mean_TYPE_A_segment_WC_reconstruct_loss = torch.mean(torch.stack(SUPER_ALL_TYPE_A_segment_WC_reconstruct_loss)) | |
| if self.TYPE_B: | |
| mean_TYPE_B_segment_termination_loss = torch.mean(torch.stack(SUPER_ALL_TYPE_B_segment_termination_loss)) | |
| mean_TYPE_B_segment_loc_reconstruct_loss = torch.mean(torch.stack(SUPER_ALL_TYPE_B_segment_loc_reconstruct_loss)) | |
| mean_TYPE_B_segment_touch_reconstruct_loss = torch.mean(torch.stack(SUPER_ALL_TYPE_B_segment_touch_reconstruct_loss)) | |
| if self.REC: | |
| mean_TYPE_B_segment_WC_reconstruct_loss = torch.mean(torch.stack(SUPER_ALL_TYPE_B_segment_WC_reconstruct_loss)) | |
| total_segment_loss = mean_segment_W_consistency_loss + mean_ORIGINAL_segment_termination_loss + mean_ORIGINAL_segment_loc_reconstruct_loss + mean_ORIGINAL_segment_touch_reconstruct_loss + mean_TYPE_A_segment_termination_loss + mean_TYPE_A_segment_loc_reconstruct_loss + mean_TYPE_A_segment_touch_reconstruct_loss + mean_TYPE_B_segment_termination_loss + mean_TYPE_B_segment_loc_reconstruct_loss + mean_TYPE_B_segment_touch_reconstruct_loss + mean_TYPE_A_segment_WC_reconstruct_loss + mean_TYPE_B_segment_WC_reconstruct_loss | |
| segment_losses = [total_segment_loss, mean_segment_W_consistency_loss, mean_ORIGINAL_segment_termination_loss, mean_ORIGINAL_segment_loc_reconstruct_loss, mean_ORIGINAL_segment_touch_reconstruct_loss, mean_TYPE_A_segment_termination_loss, mean_TYPE_A_segment_loc_reconstruct_loss, mean_TYPE_A_segment_touch_reconstruct_loss, mean_TYPE_B_segment_termination_loss, mean_TYPE_B_segment_loc_reconstruct_loss, mean_TYPE_B_segment_touch_reconstruct_loss, mean_TYPE_A_segment_WC_reconstruct_loss, mean_TYPE_B_segment_WC_reconstruct_loss] | |
| total_loss = total_sentence_loss + total_word_loss + total_segment_loss | |
| return total_loss, sentence_losses, word_losses, segment_losses | |
| def sample(self, inputs): | |
| [ word_level_stroke_in, word_level_stroke_out, word_level_stroke_length, | |
| word_level_term, word_level_char, word_level_char_length, segment_level_stroke_in, | |
| segment_level_stroke_out, segment_level_stroke_length, segment_level_term, | |
| segment_level_char, segment_level_char_length ] = inputs | |
| word_inf_state_out = self.inf_state_fc1(word_level_stroke_out[0]) | |
| word_inf_state_out = self.inf_state_relu(word_inf_state_out) | |
| word_inf_state_out, (c,h) = self.inf_state_lstm(word_inf_state_out) | |
| user_word_level_char = word_level_char[0] | |
| user_word_level_term = word_level_term[0] | |
| raw_Ws = [] | |
| original_Wc = [] | |
| word_batch_id = 0 | |
| # ORIGINAL | |
| curr_seq_len = word_level_stroke_length[0][word_batch_id][0] | |
| curr_char_len = word_level_char_length[0][word_batch_id][0] | |
| char_vector = torch.eye(len(CHARACTERS))[user_word_level_char[word_batch_id][:curr_char_len]].to(self.device) | |
| current_term = user_word_level_term[word_batch_id][:curr_seq_len].unsqueeze(-1) | |
| split_ids = torch.nonzero(current_term)[:,0] | |
| # char_vector = self.char_vec_fc(char_vector) | |
| # char_vector = self.char_vec_relu(char_vector) | |
| char_vector_1 = self.char_vec_fc_1(char_vector) | |
| char_vector_1 = self.char_vec_relu_1(char_vector_1) | |
| # unique_char_matrices = [] | |
| # for cid in range(len(char_vector)): | |
| # unique_char_vector = char_vector[cid:cid+1] | |
| # unique_char_out = unique_char_vector.unsqueeze(0) | |
| # unique_char_out, (c,h) = self.char_lstm(unique_char_out) | |
| # unique_char_out = unique_char_out.squeeze(0) | |
| # unique_char_out = self.char_vec_fc2(unique_char_out) | |
| # unique_char_matrix = unique_char_out.view([-1,1,self.weight_dim,self.weight_dim]) | |
| # unique_char_matrix = unique_char_matrix.squeeze(1) | |
| # unique_char_matrices.append(unique_char_matrix) | |
| unique_char_matrices_1 = [] | |
| for cid in range(len(char_vector)): | |
| # Tower 1 | |
| unique_char_vector_1 = char_vector_1[cid:cid+1] | |
| unique_char_input_1 = unique_char_vector_1.unsqueeze(0) | |
| unique_char_out_1, (c,h) = self.char_lstm_1(unique_char_input_1) | |
| unique_char_out_1 = unique_char_out_1.squeeze(0) | |
| unique_char_out_1 = self.char_vec_fc2_1(unique_char_out_1) | |
| unique_char_matrix_1 = unique_char_out_1.view([-1,1,self.weight_dim,self.weight_dim]) | |
| unique_char_matrix_1 = unique_char_matrix_1.squeeze(1) | |
| unique_char_matrices_1.append(unique_char_matrix_1) | |
| # Tower 1 | |
| char_out_1 = char_vector_1.unsqueeze(0) | |
| char_out_1, (c,h) = self.char_lstm_1(char_out_1) | |
| char_out_1 = char_out_1.squeeze(0) | |
| char_out_1 = self.char_vec_fc2_1(char_out_1) | |
| char_matrix_1 = char_out_1.view([-1,1,self.weight_dim,self.weight_dim]) | |
| char_matrix_1 = char_matrix_1.squeeze(1) | |
| char_matrix_inv_1 = torch.inverse(char_matrix_1) | |
| W_c_t = word_inf_state_out[word_batch_id][:curr_seq_len] | |
| W_c = torch.stack([W_c_t[i] for i in split_ids]) | |
| original_Wc.append(W_c) | |
| W = torch.bmm(char_matrix_inv_1, | |
| W_c.unsqueeze(2)).squeeze(-1) | |
| user_segment_level_stroke_length = segment_level_stroke_length[0][word_batch_id] | |
| user_segment_level_char_length = segment_level_char_length[0][word_batch_id] | |
| user_segment_level_term = segment_level_term[0][word_batch_id] | |
| user_segment_level_char = segment_level_char[0][word_batch_id] | |
| user_segment_level_stroke_in = segment_level_stroke_in[0][word_batch_id] | |
| user_segment_level_stroke_out = segment_level_stroke_out[0][word_batch_id] | |
| segment_inf_state_out = self.inf_state_fc1(user_segment_level_stroke_out) | |
| segment_inf_state_out = self.inf_state_relu(segment_inf_state_out) | |
| segment_inf_state_out, (c,h) = self.inf_state_lstm(segment_inf_state_out) | |
| segment_W_c = [] | |
| for segment_batch_id in range(len(user_segment_level_char)): | |
| curr_seq_len = user_segment_level_stroke_length[segment_batch_id][0] | |
| curr_char_len = user_segment_level_char_length[segment_batch_id][0] | |
| current_term = user_segment_level_term[segment_batch_id][:curr_seq_len].unsqueeze(-1) | |
| split_ids = torch.nonzero(current_term)[:,0] | |
| seg_W_c_t = segment_inf_state_out[segment_batch_id][:curr_seq_len] | |
| seg_W_c = torch.stack([seg_W_c_t[i] for i in split_ids]) | |
| segment_W_c.append(seg_W_c) | |
| target_characters_ids = word_level_char[0][0][:word_level_char_length[0][0]] | |
| target_characters = ''.join([CHARACTERS[i] for i in target_characters_ids]) | |
| mean_global_W = torch.mean(W, 0) | |
| TYPE_A_WC = torch.bmm(char_matrix_1, | |
| mean_global_W.repeat(char_matrix_1.size(0), 1).unsqueeze(2)).squeeze(-1) | |
| unique_char_matrix_1 = torch.stack(unique_char_matrices_1) | |
| unique_char_matrix_1 = unique_char_matrix_1.squeeze(1) | |
| TYPE_B_WC_RAW = torch.bmm(unique_char_matrix_1, | |
| mean_global_W.repeat(unique_char_matrix_1.size(0), 1).unsqueeze(2)).squeeze(-1) | |
| TYPE_B_WC_RAW = TYPE_B_WC_RAW.unsqueeze(0) | |
| TYPE_B_WC, (c,h) = self.magic_lstm(TYPE_B_WC_RAW) | |
| TYPE_B_WC = TYPE_B_WC.squeeze(0) | |
| # CC | |
| TYPE_C_WC = [] | |
| for segment_batch_id in range(len(segment_W_c)): | |
| if segment_batch_id == 0: | |
| for each_segment_Wc in segment_W_c[segment_batch_id]: | |
| TYPE_C_WC.append(each_segment_Wc) | |
| prev_id = len(TYPE_C_WC) - 1 | |
| else: | |
| prev_original_W_c = W_c[prev_id] | |
| for each_segment_Wc in segment_W_c[segment_batch_id]: | |
| magic_inp = torch.stack([prev_original_W_c, each_segment_Wc]) | |
| magic_inp = magic_inp.unsqueeze(0) | |
| type_c_out, (c,h) = self.magic_lstm(magic_inp) | |
| type_c_out = type_c_out.squeeze(0) | |
| TYPE_C_WC.append(type_c_out[-1]) | |
| prev_id = len(TYPE_C_WC) - 1 | |
| TYPE_C_WC = torch.stack(TYPE_C_WC) | |
| # DD | |
| TYPE_D_WC = [] | |
| TYPE_D_REF = [] | |
| for segment_batch_id in range(len(segment_W_c)): | |
| if segment_batch_id == 0: | |
| for each_segment_Wc in segment_W_c[segment_batch_id]: | |
| TYPE_D_WC.append(each_segment_Wc) | |
| TYPE_D_REF.append(segment_W_c[segment_batch_id][-1]) | |
| else: | |
| for each_segment_Wc in segment_W_c[segment_batch_id]: | |
| magic_inp = torch.cat([torch.stack(TYPE_D_REF, 0), each_segment_Wc.unsqueeze(0)], 0) | |
| magic_inp = magic_inp.unsqueeze(0) | |
| TYPE_D_out, (c,h) = self.magic_lstm(magic_inp) | |
| TYPE_D_out = TYPE_D_out.squeeze(0) | |
| TYPE_D_WC.append(TYPE_D_out[-1]) | |
| TYPE_D_REF.append(segment_W_c[segment_batch_id][-1]) | |
| TYPE_D_WC = torch.stack(TYPE_D_WC) | |
| o_tc = ''.join([CHARACTERS[c] for c in word_level_char[0][0][:word_level_char_length[0][0]]]) | |
| o_commands = self.sample_from_w(original_Wc[0], o_tc) | |
| if len(TYPE_A_WC) == len(original_Wc[0]): | |
| a_commands = self.sample_from_w(TYPE_A_WC, target_characters) | |
| else: | |
| a_commands = [[0,0,0]] | |
| if len(TYPE_B_WC) == len(original_Wc[0]): | |
| b_commands = self.sample_from_w(TYPE_B_WC, target_characters) | |
| else: | |
| b_commands = [[0,0,0]] | |
| if len(TYPE_C_WC) == len(original_Wc[0]): | |
| c_commands = self.sample_from_w(TYPE_C_WC, target_characters) | |
| else: | |
| c_commands = [[0,0,0]] | |
| if len(TYPE_D_WC) == len(original_Wc[0]): | |
| d_commands = self.sample_from_w(TYPE_D_WC, target_characters) | |
| else: | |
| d_commands = [[0,0,0]] | |
| return [word_level_stroke_out[0][0], o_commands, a_commands, b_commands, c_commands, d_commands] | |
| def sample_from_w(self, W_c_rec, target_sentence): | |
| gen_input = torch.zeros([1, 1, 3]).to(self.device) | |
| current_char_id_count = 0 | |
| gc1 = torch.zeros([self.num_layers, 1, self.weight_dim]).to(self.device) | |
| gh1 = torch.zeros([self.num_layers, 1, self.weight_dim]).to(self.device) | |
| gc2 = torch.zeros([self.num_layers, 1, self.weight_dim * 2]).to(self.device) | |
| gh2 = torch.zeros([self.num_layers, 1, self.weight_dim * 2]).to(self.device) | |
| terms = [] | |
| commands = [] | |
| character_nums = 0 | |
| cx, cy = 100, 150 | |
| for zz in range(800): | |
| W_c_t_now = W_c_rec[current_char_id_count:current_char_id_count + 1] | |
| gen_state = self.gen_state_fc1(gen_input) | |
| gen_state = self.gen_state_relu(gen_state) | |
| gen_state, (gc1, gh1) = self.gen_state_lstm1(gen_state, (gc1, gh1)) | |
| gen_encoded = gen_state.squeeze(0) | |
| gen_lstm2_input = torch.cat([gen_encoded, W_c_t_now], -1) | |
| gen_lstm2_input = gen_lstm2_input.view([1, 1, self.weight_dim * 2]) | |
| gen_out, (gc2, gh2) = self.gen_state_lstm2(gen_lstm2_input, (gc2, gh2)) | |
| gen_out = gen_out.squeeze(0) | |
| mdn_out = self.gen_state_fc2(gen_out) | |
| term_out = self.term_fc1(gen_out) | |
| term_out = self.term_relu1(term_out) | |
| term_out = self.term_fc2(term_out) | |
| term_out = self.term_relu2(term_out) | |
| term_out = self.term_fc3(term_out) | |
| term = self.term_sigmoid(term_out) | |
| eos = self.mdn_sigmoid(mdn_out[:, 0]) | |
| [mu1, mu2, sig1, sig2, rho, pi] = torch.split(mdn_out[:, 1:], self.num_mixtures, 1) | |
| sig1 = sig1.exp() + 1e-3 | |
| sig2 = sig2.exp() + 1e-3 | |
| rho = self.mdn_tanh(rho) | |
| pi = self.mdn_softmax(pi) | |
| mus = torch.stack([mu1, mu2], -1).squeeze() | |
| pi = pi.cpu().detach().numpy() | |
| mus = mus.cpu().detach().numpy() | |
| rho = rho.cpu().detach().numpy()[0] | |
| eos = eos.cpu().detach().numpy()[0] | |
| term = term.cpu().detach().numpy()[0][0] | |
| terms.append(term) | |
| [dx, dy] = np.sum(pi.reshape(20, 1) * mus, 0) | |
| # print (eos) | |
| touch = 1 if eos > 0.5 else 0 | |
| commands.append([dx, dy, touch]) | |
| gen_input = torch.FloatTensor([dx, dy, touch]).view([1, 1, 3]).to(self.device) | |
| character_nums += 1 | |
| # print (zz, term) | |
| if term > 0.3: | |
| if target_sentence[current_char_id_count] == ' ': | |
| current_char_id_count += 1 | |
| character_nums = 0 | |
| if current_char_id_count == len(W_c_rec): | |
| break | |
| elif character_nums > 5: | |
| current_char_id_count += 1 | |
| character_nums = 0 | |
| if current_char_id_count == len(W_c_rec): | |
| break | |
| cx += dx * 2.0 * 5.0 | |
| cy += dy * 2.0 * 5.0 | |
| if cx > 1000 or cx < 0: | |
| break | |
| if cy > 350 or cy < 0: | |
| break | |
| return commands | |
| def sample_from_w_fix(self, W_c_rec): | |
| gen_input = torch.zeros([1, 1, 3]).to(self.device) | |
| current_char_id_count = 0 | |
| gc1 = torch.zeros([self.num_layers, 1, self.weight_dim]).to(self.device) | |
| gh1 = torch.zeros([self.num_layers, 1, self.weight_dim]).to(self.device) | |
| gc2 = torch.zeros([self.num_layers, 1, self.weight_dim * 2]).to(self.device) | |
| gh2 = torch.zeros([self.num_layers, 1, self.weight_dim * 2]).to(self.device) | |
| terms = [] | |
| commands = [] | |
| character_nums = 0 | |
| cx, cy = 100, 150 | |
| new_char = False | |
| renewal = False | |
| for zz in range(800): | |
| # print (torch.sum(gc1)) | |
| W_c_t_now = W_c_rec[current_char_id_count:current_char_id_count + 1] | |
| gen_state = self.gen_state_fc1(gen_input) | |
| gen_state = self.gen_state_relu(gen_state) | |
| gen_state, (gc1, gh1) = self.gen_state_lstm1(gen_state, (gc1, gh1)) | |
| gen_encoded = gen_state.squeeze(0) | |
| gen_lstm2_input = torch.cat([gen_encoded, W_c_t_now], -1) | |
| gen_lstm2_input = gen_lstm2_input.view([1, 1, self.weight_dim * 2]) | |
| gen_out, (gc2, gh2) = self.gen_state_lstm2(gen_lstm2_input, (gc2, gh2)) | |
| gen_out = gen_out.squeeze(0) | |
| mdn_out = self.gen_state_fc2(gen_out) | |
| term_out = self.term_fc1(gen_out) | |
| term_out = self.term_relu1(term_out) | |
| term_out = self.term_fc2(term_out) | |
| term_out = self.term_relu2(term_out) | |
| term_out = self.term_fc3(term_out) | |
| term = self.term_sigmoid(term_out) | |
| eos = self.mdn_sigmoid(mdn_out[:, 0]) | |
| [mu1, mu2, sig1, sig2, rho, pi] = torch.split(mdn_out[:, 1:], self.num_mixtures, 1) | |
| sig1 = sig1.exp() + 1e-3 | |
| sig2 = sig2.exp() + 1e-3 | |
| rho = self.mdn_tanh(rho) | |
| pi = self.mdn_softmax(pi) | |
| mus = torch.stack([mu1, mu2], -1).squeeze() | |
| sigs = torch.stack([sig1, sig2], -1).squeeze() * self.scale_sd | |
| distribution = torch.distributions.normal.Normal(loc=mus, scale=sigs) | |
| sample = distribution.sample() | |
| min_clamp = distribution.icdf(0.5 - torch.ones_like(mus) * self.clamp_mdn/2) | |
| max_clamp = distribution.icdf(0.5 + torch.ones_like(mus) * self.clamp_mdn/2) | |
| sample = sample.clamp(min=min_clamp, max=max_clamp) | |
| pi = pi.cpu().detach().numpy() | |
| mus = mus.cpu().detach().numpy() | |
| rho = rho.cpu().detach().numpy()[0] | |
| eos = eos.cpu().detach().numpy()[0] | |
| term = term.cpu().detach().numpy()[0][0] | |
| sample = sample.cpu().detach().numpy() | |
| terms.append(term) | |
| [dx, dy] = np.sum(pi.reshape(20, 1) * sample, 0) | |
| touch = 1 if eos > 0.5 else 0 | |
| if new_char and touch == 1: | |
| new_char = False | |
| commands.append([dx, dy, touch]) | |
| return commands, current_char_id_count | |
| else: | |
| commands.append([dx, dy, touch]) | |
| gen_input = torch.FloatTensor([dx, dy, touch]).view([1, 1, 3]).to(self.device) | |
| character_nums += 1 | |
| # print (zz, term) | |
| if term > 0.5: | |
| if character_nums > 5: | |
| current_char_id_count += 1 | |
| character_nums = 0 | |
| new_char = True | |
| if current_char_id_count == len(W_c_rec): | |
| break | |
| cx += dx * 2.0 * 5.0 | |
| cy += dy * 2.0 * 5.0 | |
| if cx > 1000 or cx < 0: | |
| break | |
| if cy > 350 or cy < 0: | |
| break | |
| return commands, -1 |