def main():
print('go to model')
print '*' * 80
spk_global_gen=prepare_data(mode='global',train_or_test='train') #写一个假的数据生成，可以用来写模型先
global_para=spk_global_gen.next()
print global_para
spk_all_list,dict_spk2idx,dict_idx2spk,mix_speech_len,speech_fre,total_frames,spk_num_total=global_para
del spk_global_gen
num_labels=len(spk_all_list)
# data_generator=prepare_data('once','train')
# data_generator=prepare_data_fake(train_or_test='train',num_labels=num_labels) #写一个假的数据生成，可以用来写模型先
#此处顺序是 mix_speechs.shape,mix_feas.shape,aim_fea.shape,aim_spkid.shape,query.shape
#一个例子：(5, 17040) (5, 134, 129) (5, 134, 129) (5,) (5, 32, 400, 300, 3)
# datasize=prepare_datasize(data_generator)
# mix_speech_len,speech_fre,total_frames,spk_num_total,video_size=datasize
print 'Begin to build the maim model for Multi_Modal Cocktail Problem.'
# data=data_generator.next()
# This part is to build the 3D mix speech embedding maps.
mix_hidden_layer_3d=MIX_SPEECH(speech_fre,mix_speech_len).cuda()
mix_speech_classifier=MIX_SPEECH_classifier(speech_fre,mix_speech_len,num_labels).cuda()
mix_speech_multiEmbedding=SPEECH_EMBEDDING(num_labels,config.EMBEDDING_SIZE,spk_num_total+config.UNK_SPK_SUPP).cuda()
print mix_hidden_layer_3d
print mix_speech_classifier
# mix_speech_hidden=mix_hidden_layer_3d(Variable(torch.from_numpy(data[1])).cuda())
# mix_speech_output=mix_speech_classifier(Variable(torch.from_numpy(data[1])).cuda())
# 技巧：alpha0的时候，就是选出top_k，top_k很大的时候，就是选出来大于alpha的
# top_k_mask_mixspeech=top_k_mask(mix_speech_output,alpha=config.ALPHA,top_k=config.MAX_MIX)
# top_k_mask_mixspeech=top_k_mask(mix_speech_output,alpha=config.ALPHA,top_k=3)
# print top_k_mask_mixspeech
# mix_speech_multiEmbs=mix_speech_multiEmbedding(top_k_mask_mixspeech) # bs*num_labels（最多混合人个数）×Embedding的大小
# mix_speech_multiEmbs=mix_speech_multiEmbedding(Variable(torch.from_numpy(top_k_mask_mixspeech),requires_grad=False).cuda()) # bs*num_labels（最多混合人个数）×Embedding的大小
# 需要计算：mix_speech_hidden[bs,len,fre,emb]和mix_mulEmbedding[bs,num_labels,EMB]的Ａttention
# 把 前者扩充为bs*num_labels,XXXXXXXXX的，后者也是，然后用ＡＴＴ函数计算它们再转回来就好了
# mix_speech_hidden_5d=mix_speech_hidden.view(config.BATCH_SIZE,1,mix_speech_len,speech_fre,config.EMBEDDING_SIZE)
# mix_speech_hidden_5d=mix_speech_hidden_5d.expand(config.BATCH_SIZE,num_labels,mix_speech_len,speech_fre,config.EMBEDDING_SIZE).contiguous()
# mix_speech_hidden_5d=mix_speech_hidden_5d.view(-1,mix_speech_len,speech_fre,config.EMBEDDING_SIZE)
# att_speech_layer=ATTENTION(config.EMBEDDING_SIZE,'align').cuda()
# att_multi_speech=att_speech_layer(mix_speech_hidden_5d,mix_speech_multiEmbs.view(-1,config.EMBEDDING_SIZE))
# print att_multi_speech.size()
# att_multi_speech=att_multi_speech.view(config.BATCH_SIZE,num_labels,mix_speech_len,speech_fre,-1)
# print att_multi_speech.size()
# This part is to conduct the video inputs.
# query_video_layer=VIDEO_QUERY(total_frames,config.VideoSize,spk_num_total).cuda()
query_video_layer=None
# print query_video_layer
# query_video_output,xx=query_video_layer(Variable(torch.from_numpy(data[4])))
# This part is to conduct the memory.
# hidden_size=(config.HIDDEN_UNITS)
hidden_size=(config.EMBEDDING_SIZE)
# x=torch.arange(0,24).view(2,3,4)
# y=torch.ones([2,4])
att_layer=ATTENTION(config.EMBEDDING_SIZE,'align').cuda()
att_speech_layer=ATTENTION(config.EMBEDDING_SIZE,'align').cuda()
# att=ATTENTION(4,'align')
# mask=att(x,y)#bs*max_len
# del data_generator
# del data
optimizer = torch.optim.Adam([{'params':mix_hidden_layer_3d.parameters()},
{'params':mix_speech_multiEmbedding.parameters()},
{'params':mix_speech_classifier.parameters()},
# {'params':query_video_layer.lstm_layer.parameters()},
# {'params':query_video_layer.dense.parameters()},
# {'params':query_video_layer.Linear.parameters()},
{'params':att_layer.parameters()},
{'params':att_speech_layer.parameters()},
# ], lr=0.02,momentum=0.9)
], lr=0.0002)
if 0 and config.Load_param:
# query_video_layer.load_state_dict(torch.load('param_video_layer_19'))
# mix_speech_classifier.load_state_dict(torch.load('params/param_speech_multilabel_epoch249'))
mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix101_WSJ0_hidden3d_180'))
mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix101_WSJ0_emblayer_180'))
att_speech_layer.load_state_dict(torch.load('params/param_mix101_WSJ0_attlayer_180'))
loss_func = torch.nn.MSELoss() # the target label is NOT an one-hotted
loss_multi_func = torch.nn.MSELoss() # the target label is NOT an one-hotted
# loss_multi_func = torch.nn.L1Loss() # the target label is NOT an one-hotted
loss_query_class=torch.nn.CrossEntropyLoss()
print '''Begin to calculate.'''
for epoch_idx in range(config.MAX_EPOCH):
if epoch_idx%50==0:
for ee in optimizer.param_groups:
ee['lr']/=2
if epoch_idx>0:
print 'SDR_SUM (len:{}) for epoch {} : '.format(SDR_SUM.shape,epoch_idx-1,SDR_SUM.mean())
SDR_SUM=np.array([])
# print_memory_state(memory.memory)
print 'SDR_SUM for epoch {}:{}'.format(epoch_idx - 1, SDR_SUM.mean())
for batch_idx in range(config.EPOCH_SIZE):
print '*' * 40,epoch_idx,batch_idx,'*'*40
train_data_gen=prepare_data('once','train')
# train_data_gen=prepare_data('once','test')
train_data=train_data_gen.next()
'''混合语音len,fre,Emb 3D表示层'''
mix_speech_hidden=mix_hidden_layer_3d(Variable(torch.from_numpy(train_data['mix_feas'])).cuda())
# 暂时关掉video部分,因为s2 s3 s4 的视频数据不全暂时
'''Speech self Sepration 语音自分离部分'''
mix_speech_output=mix_speech_classifier(Variable(torch.from_numpy(train_data['mix_feas'])).cuda())
#从数据里得到ground truth的说话人名字和vector
# y_spk_list=[one.keys() for one in train_data['multi_spk_fea_list']]
y_spk_list= train_data['multi_spk_fea_list']
y_spk_gtruth,y_map_gtruth=multi_label_vector(y_spk_list,dict_spk2idx)
# 如果训练阶段使用Ground truth的分离结果作为判别
if config.Ground_truth:
mix_speech_output=Variable(torch.from_numpy(y_map_gtruth)).cuda()
if test_all_outputchannel: #把输入的mask改成全１，可以用来测试输出所有的channel
mix_speech_output=Variable(torch.ones(config.BATCH_SIZE,num_labels,))
y_map_gtruth=np.ones([config.BATCH_SIZE,num_labels])
top_k_mask_mixspeech=top_k_mask(mix_speech_output,alpha=0.5,top_k=num_labels) #torch.Float型的
top_k_mask_idx=[np.where(line==1)[0] for line in top_k_mask_mixspeech.numpy()]
mix_speech_multiEmbs=mix_speech_multiEmbedding(top_k_mask_mixspeech,top_k_mask_idx) # bs*num_labels（最多混合人个数）×Embedding的大小
assert len(top_k_mask_idx[0])==len(top_k_mask_idx[-1])
top_k_num=len(top_k_mask_idx[0])
#需要计算：mix_speech_hidden[bs,len,fre,emb]和mix_mulEmbedding[bs,num_labels,EMB]的Ａttention
#把 前者扩充为bs*num_labels,XXXXXXXXX的，后者也是，然后用ＡＴＴ函数计算它们再转回来就好了
mix_speech_hidden_5d=mix_speech_hidden.view(config.BATCH_SIZE,1,mix_speech_len,speech_fre,config.EMBEDDING_SIZE)
mix_speech_hidden_5d=mix_speech_hidden_5d.expand(config.BATCH_SIZE,top_k_num,mix_speech_len,speech_fre,config.EMBEDDING_SIZE).contiguous()
mix_speech_hidden_5d_last=mix_speech_hidden_5d.view(-1,mix_speech_len,speech_fre,config.EMBEDDING_SIZE)
# att_speech_layer=ATTENTION(config.EMBEDDING_SIZE,'align').cuda()
att_speech_layer=ATTENTION(config.EMBEDDING_SIZE,'dot').cuda()
att_multi_speech=att_speech_layer(mix_speech_hidden_5d_last,mix_speech_multiEmbs.view(-1,config.EMBEDDING_SIZE))
# print att_multi_speech.size()
att_multi_speech=att_multi_speech.view(config.BATCH_SIZE,top_k_num,mix_speech_len,speech_fre) # bs,num_labels,len,fre这个东西
# print att_multi_speech.size()
multi_mask=att_multi_speech
# top_k_mask_mixspeech_multi=top_k_mask_mixspeech.view(config.BATCH_SIZE,top_k_num,1,1).expand(config.BATCH_SIZE,top_k_num,mix_speech_len,speech_fre)
# multi_mask=multi_mask*Variable(top_k_mask_mixspeech_multi).cuda()
x_input_map=Variable(torch.from_numpy(train_data['mix_feas'])).cuda()
# print x_input_map.size()
x_input_map_multi=x_input_map.view(config.BATCH_SIZE,1,mix_speech_len,speech_fre).expand(config.BATCH_SIZE,top_k_num,mix_speech_len,speech_fre)
# predict_multi_map=multi_mask*x_input_map_multi
predict_multi_map=multi_mask*x_input_map_multi
if 0 and batch_idx%100==0:
print multi_mask
# print predict_multi_map
y_multi_map=np.zeros([config.BATCH_SIZE,top_k_num,mix_speech_len,speech_fre],dtype=np.float32)
batch_spk_multi_dict=train_data['multi_spk_fea_list']
for idx,sample in enumerate(batch_spk_multi_dict):
y_idx=sorted([dict_spk2idx[spk] for spk in sample.keys()])
assert y_idx==list(top_k_mask_idx[idx])
for jdx,oo in enumerate(y_idx):
y_multi_map[idx,jdx]=sample[dict_idx2spk[oo]]
y_multi_map= Variable(torch.from_numpy(y_multi_map)).cuda()
loss_multi_speech=loss_multi_func(predict_multi_map,y_multi_map)
#各通道和为１的loss部分,应该可以更多的带来差异
y_sum_map=Variable(torch.ones(config.BATCH_SIZE,mix_speech_len,speech_fre)).cuda()
predict_sum_map=torch.sum(multi_mask,1)
loss_multi_sum_speech=loss_multi_func(predict_sum_map,y_sum_map)
# loss_multi_speech=loss_multi_speech #todo:以后可以研究下这个和为１的效果对比一下，暂时直接MSE效果已经很不错了。
print 'loss 1, losssum : ',loss_multi_speech.data.cpu().numpy(),loss_multi_sum_speech.data.cpu().numpy()
loss_multi_speech=loss_multi_speech+0.5*loss_multi_sum_speech
print 'training multi-abs norm this batch:',torch.abs(y_multi_map-predict_multi_map).norm().data.cpu().numpy()
print 'loss:',loss_multi_speech.data.cpu().numpy()
if 1 or batch_idx==config.EPOCH_SIZE-1:
bss_eval(predict_multi_map,y_multi_map,top_k_mask_idx,dict_idx2spk,train_data)
SDR_SUM = np.append(SDR_SUM, bss_test.cal('batch_output/', 2))
optimizer.zero_grad() # clear gradients for next train
loss_multi_speech.backward() # backpropagation, compute gradients
optimizer.step() # apply gradients
if 1 and epoch_idx>80 and epoch_idx%10==0 and batch_idx==config.EPOCH_SIZE-1:
torch.save(mix_speech_multiEmbedding.state_dict(),'params/param_mix_{}_emblayer_{}'.format(config.DATASET,epoch_idx))
torch.save(mix_hidden_layer_3d.state_dict(),'params/param_mix_{}_hidden3d_{}'.format(config.DATASET,epoch_idx))
torch.save(att_speech_layer.state_dict(),'params/param_mix_{}_attlayer_{}'.format(config.DATASET,epoch_idx))
def main():
print('go to model')
print '*' * 80
spk_global_gen = prepare_data(mode='global',
train_or_test='train') #写一个假的数据生成，可以用来写模型先
global_para = spk_global_gen.next()
print global_para
spk_all_list, dict_spk2idx, dict_idx2spk, mix_speech_len, speech_fre, total_frames, spk_num_total = global_para
del spk_global_gen
num_labels = len(spk_all_list)
# data_generator=prepare_data('once','train')
# data_generator=prepare_data_fake(train_or_test='train',num_labels=num_labels) #写一个假的数据生成，可以用来写模型先
#此处顺序是 mix_speechs.shape,mix_feas.shape,aim_fea.shape,aim_spkid.shape,query.shape
#一个例子：(5, 17040) (5, 134, 129) (5, 134, 129) (5,) (5, 32, 400, 300, 3)
# datasize=prepare_datasize(data_generator)
# mix_speech_len,speech_fre,total_frames,spk_num_total,video_size=datasize
print 'Begin to build the maim model for Multi_Modal Cocktail Problem.'
# data=data_generator.next()
# This part is to build the 3D mix speech embedding maps.
mix_hidden_layer_3d = MIX_SPEECH(speech_fre, mix_speech_len).cuda()
mix_speech_classifier = MIX_SPEECH_classifier(speech_fre, mix_speech_len,
num_labels).cuda()
mix_speech_multiEmbedding = SPEECH_EMBEDDING(
num_labels, config.EMBEDDING_SIZE,
spk_num_total + config.UNK_SPK_SUPP).cuda()
print mix_hidden_layer_3d
print mix_speech_classifier
# mix_speech_hidden=mix_hidden_layer_3d(Variable(torch.from_numpy(data[1])).cuda())
# mix_speech_output=mix_speech_classifier(Variable(torch.from_numpy(data[1])).cuda())
# 技巧：alpha0的时候，就是选出top_k，top_k很大的时候，就是选出来大于alpha的
# top_k_mask_mixspeech=top_k_mask(mix_speech_output,alpha=config.ALPHA,top_k=config.MAX_MIX)
# top_k_mask_mixspeech=top_k_mask(mix_speech_output,alpha=config.ALPHA,top_k=3)
# print top_k_mask_mixspeech
# mix_speech_multiEmbs=mix_speech_multiEmbedding(top_k_mask_mixspeech) # bs*num_labels（最多混合人个数）×Embedding的大小
# mix_speech_multiEmbs=mix_speech_multiEmbedding(Variable(torch.from_numpy(top_k_mask_mixspeech),requires_grad=False).cuda()) # bs*num_labels（最多混合人个数）×Embedding的大小
# 需要计算：mix_speech_hidden[bs,len,fre,emb]和mix_mulEmbedding[bs,num_labels,EMB]的Ａttention
# 把 前者扩充为bs*num_labels,XXXXXXXXX的，后者也是，然后用ＡＴＴ函数计算它们再转回来就好了
# mix_speech_hidden_5d=mix_speech_hidden.view(config.BATCH_SIZE,1,mix_speech_len,speech_fre,config.EMBEDDING_SIZE)
# mix_speech_hidden_5d=mix_speech_hidden_5d.expand(config.BATCH_SIZE,num_labels,mix_speech_len,speech_fre,config.EMBEDDING_SIZE).contiguous()
# mix_speech_hidden_5d=mix_speech_hidden_5d.view(-1,mix_speech_len,speech_fre,config.EMBEDDING_SIZE)
# att_speech_layer=ATTENTION(config.EMBEDDING_SIZE,'align').cuda()
# att_multi_speech=att_speech_layer(mix_speech_hidden_5d,mix_speech_multiEmbs.view(-1,config.EMBEDDING_SIZE))
# print att_multi_speech.size()
# att_multi_speech=att_multi_speech.view(config.BATCH_SIZE,num_labels,mix_speech_len,speech_fre,-1)
# print att_multi_speech.size()
# This part is to conduct the video inputs.
# query_video_layer=VIDEO_QUERY(total_frames,config.VideoSize,spk_num_total).cuda()
query_video_layer = None
# print query_video_layer
# query_video_output,xx=query_video_layer(Variable(torch.from_numpy(data[4])))
# This part is to conduct the memory.
# hidden_size=(config.HIDDEN_UNITS)
hidden_size = (config.EMBEDDING_SIZE)
# x=torch.arange(0,24).view(2,3,4)
# y=torch.ones([2,4])
att_layer = ATTENTION(config.EMBEDDING_SIZE, 'align').cuda()
att_speech_layer = ATTENTION(config.EMBEDDING_SIZE, 'align').cuda()
# att=ATTENTION(4,'align')
# mask=att(x,y)#bs*max_len
# del data_generator
# del data
optimizer = torch.optim.Adam(
[
{
'params': mix_hidden_layer_3d.parameters()
},
{
'params': mix_speech_multiEmbedding.parameters()
},
{
'params': mix_speech_classifier.parameters()
},
# {'params':query_video_layer.lstm_layer.parameters()},
# {'params':query_video_layer.dense.parameters()},
# {'params':query_video_layer.Linear.parameters()},
{
'params': att_layer.parameters()
},
{
'params': att_speech_layer.parameters()
},
# ], lr=0.02,momentum=0.9)
],
lr=0.00005)
if 1 and config.Load_param:
# query_video_layer.load_state_dict(torch.load('param_video_layer_19'))
# mix_speech_classifier.load_state_dict(torch.load('params/param_speech_123onezeroag3_WSJ0_multilabel_epoch40'))
mix_speech_classifier.load_state_dict(
torch.load(
'params/param_speech_123onezeroag4_WSJ0_multilabel_epoch70'))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix101_WSJ0_hidden3d_180'))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix101_WSJ0_emblayer_180'))
# att_speech_layer.load_state_dict(torch.load('params/param_mix101_WSJ0_attlayer_180'))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix101_dbag1nosum_WSJ0_hidden3d_350',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix101_dbag1nosum_WSJ0_emblayer_350',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix101_dbag1nosum_WSJ0_attlayer_350',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix2or3_db_WSJ0_hidden3d_560',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix2or3_db_WSJ0_emblayer_560',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix2or3_db_WSJ0_attlayer_560',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix101_dbag2sum_WSJ0_hidden3d_460',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix101_dbag2sum_WSJ0_emblayer_460',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix101_dbag2sum_WSJ0_attlayer_460',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix101_dbdropout_WSJ0_hidden3d_370',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix101_dbdropout_WSJ0_emblayer_370',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix101_dbdropout_WSJ0_attlayer_370',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix101_4db_WSJ0_hidden3d_110',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix101_4db_WSJ0_emblayer_110',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix101_4db_WSJ0_attlayer_110',map_location={'cuda:1':'cuda:0'}))
# mix_speech_classifier.load_state_dict(torch.load('params/param_speech_4lstm_multilabelloss30map_epoch440'))
# att_speech_layer.load_state_dict(torch.load('params/param_mix101_dbdropoutag_WSJ0_attlayer_220',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix101_dbdropoutag_WSJ0_hidden3d_220',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix101_dbdropoutag_WSJ0_emblayer_220',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix2_db2dropout_WSJ0_attlayer_495',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix2_db2dropout_WSJ0_hidden3d_495',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix2_db2dropout_WSJ0_emblayer_495',map_location={'cuda:1':'cuda:0'}))
att_speech_layer.load_state_dict(
torch.load('params/param_mix2_db2dropout_WSJ0_attlayer_90',
map_location={'cuda:1': 'cuda:0'}))
mix_hidden_layer_3d.load_state_dict(
torch.load('params/param_mix2_db2dropout_WSJ0_hidden3d_90',
map_location={'cuda:1': 'cuda:0'}))
mix_speech_multiEmbedding.load_state_dict(
torch.load('params/param_mix2_db2dropout_WSJ0_emblayer_90',
map_location={'cuda:1': 'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix1to3_dbdropoutag1_WSJ0_attlayer_430',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix1to3_dbdropoutag1_WSJ0_hidden3d_430',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix1to3_dbdropoutag1_WSJ0_emblayer_430',map_location={'cuda:1':'cuda:0'}))
loss_func = torch.nn.MSELoss() # the target label is NOT an one-hotted
loss_multi_func = torch.nn.MSELoss(
) # the target label is NOT an one-hotted
# loss_multi_func = torch.nn.L1Loss() # the target label is NOT an one-hotted
loss_query_class = torch.nn.CrossEntropyLoss()
print '''Begin to calculate.'''
for epoch_idx in range(config.MAX_EPOCH):
if epoch_idx > 0:
print 'SDR_SUM (len:{}) for epoch {} : {}'.format(
SDR_SUM.shape, epoch_idx - 1, SDR_SUM.mean())
SDR_SUM = np.array([])
# print_memory_state(memory.memory)
print 'SDR_SUM for epoch {}:{}'.format(epoch_idx - 1, SDR_SUM.mean())
for batch_idx in range(config.EPOCH_SIZE):
print '*' * 40, epoch_idx, batch_idx, '*' * 40
train_data_gen = prepare_data('once', 'train')
# train_data_gen=prepare_data('once','test')
# train_data_gen=prepare_data('once','eval_test')
train_data = train_data_gen.next()
'''混合语音len,fre,Emb 3D表示层'''
mix_speech_hidden = mix_hidden_layer_3d(
Variable(torch.from_numpy(train_data['mix_feas'])).cuda())
# 暂时关掉video部分,因为s2 s3 s4 的视频数据不全暂时
'''Speech self Sepration 语音自分离部分'''
mix_speech_output = mix_speech_classifier(
Variable(torch.from_numpy(train_data['mix_feas'])).cuda())
#从数据里得到ground truth的说话人名字和vector
y_spk_list = [
one.keys() for one in train_data['multi_spk_fea_list']
]
y_spk_list = train_data['multi_spk_fea_list']
y_spk_gtruth, y_map_gtruth = multi_label_vector(
y_spk_list, dict_spk2idx)
# 如果训练阶段使用Ground truth的分离结果作为判别
if 1 and config.Ground_truth:
mix_speech_output = Variable(
torch.from_numpy(y_map_gtruth)).cuda()
if 0 and test_all_outputchannel: #把输入的mask改成全１，可以用来测试输出所有的channel
mix_speech_output = Variable(
torch.ones(
config.BATCH_SIZE,
num_labels,
))
y_map_gtruth = np.ones([config.BATCH_SIZE, num_labels])
max_num_labels = 2
top_k_mask_mixspeech, top_k_sort_index = top_k_mask(
mix_speech_output, alpha=-0.5,
top_k=max_num_labels) #torch.Float型的
top_k_mask_idx = [
np.where(line == 1)[0]
for line in top_k_mask_mixspeech.numpy()
]
mix_speech_multiEmbs = mix_speech_multiEmbedding(
top_k_mask_mixspeech,
top_k_mask_idx) # bs*num_labels（最多混合人个数）×Embedding的大小
assert len(top_k_mask_idx[0]) == len(top_k_mask_idx[-1])
top_k_num = len(top_k_mask_idx[0])
#需要计算：mix_speech_hidden[bs,len,fre,emb]和mix_mulEmbedding[bs,num_labels,EMB]的Ａttention
#把 前者扩充为bs*num_labels,XXXXXXXXX的，后者也是，然后用ＡＴＴ函数计算它们再转回来就好了
mix_speech_hidden_5d = mix_speech_hidden.view(
config.BATCH_SIZE, 1, mix_speech_len, speech_fre,
config.EMBEDDING_SIZE)
mix_speech_hidden_5d = mix_speech_hidden_5d.expand(
config.BATCH_SIZE, top_k_num, mix_speech_len, speech_fre,
config.EMBEDDING_SIZE).contiguous()
mix_speech_hidden_5d_last = mix_speech_hidden_5d.view(
-1, mix_speech_len, speech_fre, config.EMBEDDING_SIZE)
# att_speech_layer=ATTENTION(config.EMBEDDING_SIZE,'align').cuda()
att_speech_layer = ATTENTION(config.EMBEDDING_SIZE, 'dot').cuda()
att_multi_speech = att_speech_layer(
mix_speech_hidden_5d_last,
mix_speech_multiEmbs.view(-1, config.EMBEDDING_SIZE))
# print att_multi_speech.size()
att_multi_speech = att_multi_speech.view(
config.BATCH_SIZE, top_k_num, mix_speech_len,
speech_fre) # bs,num_labels,len,fre这个东西
# print att_multi_speech.size()
multi_mask = att_multi_speech
# top_k_mask_mixspeech_multi=top_k_mask_mixspeech.view(config.BATCH_SIZE,top_k_num,1,1).expand(config.BATCH_SIZE,top_k_num,mix_speech_len,speech_fre)
# multi_mask=multi_mask*Variable(top_k_mask_mixspeech_multi).cuda()
x_input_map = Variable(torch.from_numpy(
train_data['mix_feas'])).cuda()
# print x_input_map.size()
x_input_map_multi = x_input_map.view(
config.BATCH_SIZE, 1, mix_speech_len,
speech_fre).expand(config.BATCH_SIZE, top_k_num,
mix_speech_len, speech_fre)
# predict_multi_map=multi_mask*x_input_map_multi
predict_multi_map = multi_mask * x_input_map_multi
bss_eval_fromGenMap(multi_mask, x_input_map, top_k_mask_mixspeech,
dict_idx2spk, train_data, top_k_sort_index)
SDR_SUM = np.append(SDR_SUM, bss_test.cal('batch_output/', 2))
print 'SDR_SUM (len:{}) for epoch {} : {}'.format(
SDR_SUM.shape, epoch_idx, SDR_SUM.mean())
def main():
print('go to model')
print '*' * 80
spk_global_gen=prepare_data(mode='global',train_or_test='train') #写一个假的数据生成，可以用来写模型先
global_para=spk_global_gen.next()
print global_para
spk_all_list,dict_spk2idx,dict_idx2spk,mix_speech_len,speech_fre,total_frames,spk_num_total=global_para
del spk_global_gen
num_labels=len(spk_all_list)
#此处顺序是 mix_speechs.shape,mix_feas.shape,aim_fea.shape,aim_spkid.shape,query.shape
#一个例子：(5, 17040) (5, 134, 129) (5, 134, 129) (5,) (5, 32, 400, 300, 3)
# datasize=prepare_datasize(data_generator)
# mix_speech_len,speech_fre,total_frames,spk_num_total,video_size=datasize
print 'Begin to build the maim model for Multi_Modal Cocktail Problem.'
# data=data_generator.next()
# This part is to build the 3D mix speech embedding maps.
mix_hidden_layer_3d=MIX_SPEECH(speech_fre,mix_speech_len).cuda()
mix_speech_classifier=MIX_SPEECH_classifier(speech_fre,mix_speech_len,num_labels).cuda()
mix_speech_multiEmbedding=SPEECH_EMBEDDING(num_labels,config.EMBEDDING_SIZE,spk_num_total+config.UNK_SPK_SUPP).cuda()
print mix_hidden_layer_3d
print mix_speech_classifier
# mix_speech_hidden=mix_hidden_layer_3d(Variable(torch.from_numpy(data[1])).cuda())
hidden_size=(config.EMBEDDING_SIZE)
# x=torch.arange(0,24).view(2,3,4)
# y=torch.ones([2,4])
att_layer=ATTENTION(config.EMBEDDING_SIZE,'align').cuda()
att_speech_layer=ATTENTION(config.EMBEDDING_SIZE,'align').cuda()
print att_speech_layer
optimizer = torch.optim.Adam([{'params':mix_hidden_layer_3d.parameters()},
{'params':mix_speech_multiEmbedding.parameters()},
{'params':mix_speech_classifier.parameters()},
# {'params':query_video_layer.lstm_layer.parameters()},
# {'params':query_video_layer.dense.parameters()},
# {'params':query_video_layer.Linear.parameters()},
{'params':att_layer.parameters()},
{'params':att_speech_layer.parameters()},
# ], lr=0.02,momentum=0.9)
], lr=0.0002)
if 1 and config.Load_param:
# query_video_layer.load_state_dict(torch.load('param_video_layer_19'))
# mix_speech_classifier.load_state_dict(torch.load('params/param_speech_123onezeroag3_WSJ0_multilabel_epoch40'))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix101_WSJ0_hidden3d_180'))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix101_WSJ0_emblayer_180'))
# att_speech_layer.load_state_dict(torch.load('params/param_mix101_WSJ0_attlayer_180'))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix101_dbag1nosum_WSJ0_hidden3d_250',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix101_dbag1nosum_WSJ0_emblayer_250',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix101_dbag1nosum_WSJ0_attlayer_250',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix2or3_db_WSJ0_hidden3d_560',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix2or3_db_WSJ0_emblayer_560',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix2or3_db_WSJ0_attlayer_560',map_location={'cuda:1':'cuda:0'}))
mix_speech_classifier.load_state_dict(torch.load('params/param_speech_4lstm_multilabelloss30map_epoch440'))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix101_dbag2sum_WSJ0_hidden3d_460',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix101_dbag2sum_WSJ0_emblayer_460',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix101_dbag2sum_WSJ0_attlayer_460',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix101_dbdropout_WSJ0_hidden3d_370',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix101_dbdropout_WSJ0_emblayer_370',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix101_dbdropout_WSJ0_attlayer_370',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix101_dbdropoutag_WSJ0_attlayer_220',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix101_dbdropoutag_WSJ0_hidden3d_220',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix101_dbdropoutag_WSJ0_emblayer_220',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix2or3_dbdropoutag_WSJ0_attlayer_180',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix2or3_dbdropoutag_WSJ0_hidden3d_180',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix2or3_dbdropoutag_WSJ0_emblayer_180',map_location={'cuda:1':'cuda:0'}))
att_speech_layer.load_state_dict(torch.load('params/param_mix2_db2dropout_WSJ0_attlayer_495',map_location={'cuda:1':'cuda:0'}))
mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix2_db2dropout_WSJ0_hidden3d_495',map_location={'cuda:1':'cuda:0'}))
mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix2_db2dropout_WSJ0_emblayer_495',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix2or3_db2dropout_WSJ0_attlayer_95',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix2or3_db2dropout_WSJ0_hidden3d_95',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix2or3_db2dropout_WSJ0_emblayer_95',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix1to3_dbdropoutag1_WSJ0_attlayer_500',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix1to3_dbdropoutag1_WSJ0_hidden3d_500',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix1to3_dbdropoutag1_WSJ0_emblayer_500',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix2_40lstm2dbdro_WSJ0_attlayer_835',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix2_40lstm2dbdro_WSJ0_hidden3d_835',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix2_40lstm2dbdro_WSJ0_emblayer_835',map_location={'cuda:1':'cuda:0'}))
# att_speech_layer.load_state_dict(torch.load('params/param_mix2_40lstmdbdropout_WSJ0_attlayer_200',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix2_40lstm3dbdropout_WSJ0_hidden3d_200',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_mix2_40lstm3dbdropout_WSJ0_emblayer_200',map_location={'cuda:1':'cuda:0'}))
'''with Noise'''
# att_speech_layer.load_state_dict(torch.load('params/param_noicemix2or3_db2dropout_WSJ0_attlayer_80',map_location={'cuda:1':'cuda:0'}))
# mix_hidden_layer_3d.load_state_dict(torch.load('params/param_noicemix2or3_db2dropout_WSJ0_hidden3d_80',map_location={'cuda:1':'cuda:0'}))
# mix_speech_multiEmbedding.load_state_dict(torch.load('params/param_noicemix2or3_db2dropout_WSJ0_emblayer_80',map_location={'cuda:1':'cuda:0'}))
loss_func = torch.nn.MSELoss() # the target label is NOT an one-hotted
loss_multi_func = torch.nn.MSELoss() # the target label is NOT an one-hotted
# loss_multi_func = torch.nn.L1Loss() # the target label is NOT an one-hotted
loss_query_class=torch.nn.CrossEntropyLoss()
print '''Begin to calculate.'''
SDR_SUM_total=np.array([])
for epoch_idx in range(config.MAX_EPOCH):
if epoch_idx>0:
print 'SDR_SUM (len:{}) for epoch {} : {}'.format(SDR_SUM.shape,epoch_idx-1,SDR_SUM.mean())
SDR_SUM=np.array([])
# print 'SDR_SUM for epoch {}:{}'.format(epoch_idx - 1, SDR_SUM.mean())
dst='batch_output'
if os.path.exists(dst):
print " cleanup: " + dst + "/"
shutil.rmtree(dst)
os.makedirs(dst)
for batch_idx in range(config.EPOCH_SIZE):
print '*' * 40,epoch_idx,batch_idx,'*'*40
train_data_gen=prepare_data('once','train')
# train_data_gen=prepare_data('once','test')
train_data_gen=prepare_data('once','eval_test')
train_data=train_data_gen.next()
mix_feas=train_data['mix_feas']
'''混合语音len,fre,Emb 3D表示层'''
mix_speech_hidden=mix_hidden_layer_3d(Variable(torch.from_numpy(train_data['mix_feas'])).cuda())
# 暂时关掉video部分,因为s2 s3 s4 的视频数据不全暂时
'''Speech self Sepration 语音自分离部分'''
mix_speech_output=mix_speech_classifier(Variable(torch.from_numpy(train_data['mix_feas'])).cuda())
#从数据里得到ground truth的说话人名字和vector
# y_spk_list=[one.keys() for one in train_data['multi_spk_fea_list']]
# y_spk_list= train_data['multi_spk_fea_list']
# y_spk_gtruth,y_map_gtruth=multi_label_vector(y_spk_list,dict_spk2idx)
# 如果训练阶段使用Ground truth的分离结果作为判别
if 0 and config.Ground_truth:
mix_speech_output=Variable(torch.from_numpy(y_map_gtruth)).cuda()
if test_all_outputchannel: #把输入的mask改成全１，可以用来测试输出所有的channel
mix_speech_output=Variable(torch.ones(config.BATCH_SIZE,num_labels,))
y_map_gtruth=np.ones([config.BATCH_SIZE,num_labels])
recu_spk_list=OrderedDict() #每step对应spk以及分离出来的目标语音
speech_history=[] #将每step剩余speech 频谱的历史记录下来
bss_eval_groundtrue(train_data,batch_idx)
now_feas=train_data['mix_feas']
while True:
speech_history.append(now_feas)
max_num_labels=3
top_k_mask_mixspeech,top_k_sort_index=top_k_mask(mix_speech_output,alpha=-0.3,top_k=max_num_labels) #torch.Float型的
# top_k_mask_idx=[np.where(line==1)[0] for line in top_k_mask_mixspeech.numpy()]
top_k_mask_idx=top_k_sort_index
#过滤一下，把之前见过的spk过滤掉
print 'predict spk:',top_k_mask_idx[0]
for k in top_k_mask_idx[0]:
if k not in recu_spk_list.keys():
top_k_mask_idx=[[k]]
break
print 'flitered spk:',top_k_mask_idx[0]
# 如果过滤完了之后啥也没有了，那么就结束了
if len(top_k_mask_idx[0])==0:
break
# elif top_k_mask_idx[0][0] in speech_history
mix_speech_multiEmbs=mix_speech_multiEmbedding(top_k_mask_mixspeech,top_k_mask_idx) # bs*num_labels（最多混合人个数）×Embedding的大小
assert len(top_k_mask_idx[0])==len(top_k_mask_idx[-1])
top_k_num=len(top_k_mask_idx[0])
#需要计算：mix_speech_hidden[bs,len,fre,emb]和mix_mulEmbedding[bs,num_labels,EMB]的Ａttention
#把 前者扩充为bs*num_labels,XXXXXXXXX的，后者也是，然后用ＡＴＴ函数计算它们再转回来就好了
mix_speech_hidden_5d=mix_speech_hidden.view(config.BATCH_SIZE,1,mix_speech_len,speech_fre,config.EMBEDDING_SIZE)
mix_speech_hidden_5d=mix_speech_hidden_5d.expand(config.BATCH_SIZE,top_k_num,mix_speech_len,speech_fre,config.EMBEDDING_SIZE).contiguous()
mix_speech_hidden_5d_last=mix_speech_hidden_5d.view(-1,mix_speech_len,speech_fre,config.EMBEDDING_SIZE)
# att_speech_layer=ATTENTION(config.EMBEDDING_SIZE,'align').cuda()
att_speech_layer=ATTENTION(config.EMBEDDING_SIZE,'dot').cuda()
att_multi_speech=att_speech_layer(mix_speech_hidden_5d_last,mix_speech_multiEmbs.view(-1,config.EMBEDDING_SIZE))
# print att_multi_speech.size()
att_multi_speech=att_multi_speech.view(config.BATCH_SIZE,top_k_num,mix_speech_len,speech_fre) # bs,num_labels,len,fre这个东西
# print att_multi_speech.size()
multi_mask=att_multi_speech
# multi_mask=(att_multi_speech>0.5)
# multi_mask=Variable(torch.from_numpy(np.float32(multi_mask.data.cpu().numpy()))).cuda()
# top_k_mask_mixspeech_multi=top_k_mask_mixspeech.view(config.BATCH_SIZE,top_k_num,1,1).expand(config.BATCH_SIZE,top_k_num,mix_speech_len,speech_fre)
# multi_mask=multi_mask*Variable(top_k_mask_mixspeech_multi).cuda()
x_input_map=Variable(torch.from_numpy(now_feas)).cuda()
# print x_input_map.size()
x_input_map_multi=x_input_map.view(config.BATCH_SIZE,1,mix_speech_len,speech_fre).expand(config.BATCH_SIZE,top_k_num,mix_speech_len,speech_fre)
# predict_multi_map=multi_mask*x_input_map_multi
predict_multi_map=multi_mask*x_input_map_multi #该说话人预测出来的频谱
recu_spk_list[top_k_mask_idx[0][0]]=predict_multi_map
pre_spk=dict_idx2spk[top_k_mask_idx[0][0]]
num_step=len(recu_spk_list)
print 'Now output the {} th spk , closest to spk <{}> in train list.'.format(num_step,pre_spk)
# bss_eval_recu(multi_mask,x_input_map,top_k_mask_mixspeech,pre_spk,train_data,num_step-1,batch_idx)
if num_step>=2:
# bss_eval_recu(multi_mask,x_input_map,top_k_mask_mixspeech,pre_spk,train_data,num_step,batch_idx)
break
now_feas=((1-multi_mask)*x_input_map_multi).data.cpu().numpy().reshape(1,mix_speech_len,speech_fre)
mix_speech_output=mix_speech_classifier(Variable(torch.from_numpy(now_feas)).cuda())
mix_speech_hidden=mix_hidden_layer_3d(Variable(torch.from_numpy(now_feas)).cuda())
cal_spk=recu_spk_list.keys()
mix_speech_multiEmbs=mix_speech_multiEmbedding(top_k_mask_mixspeech,cal_spk) # bs*num_labels（最多混合人个数）×Embedding的大小
top_k_num=len(cal_spk)
#需要计算：mix_speech_hidden[bs,len,fre,emb]和mix_mulEmbedding[bs,num_labels,EMB]的Ａttention
#把 前者扩充为bs*num_labels,XXXXXXXXX的，后者也是，然后用ＡＴＴ函数计算它们再转回来就好了
mix_speech_hidden=mix_hidden_layer_3d(Variable(torch.from_numpy(train_data['mix_feas'])).cuda())
mix_speech_hidden_5d=mix_speech_hidden.view(config.BATCH_SIZE,1,mix_speech_len,speech_fre,config.EMBEDDING_SIZE)
mix_speech_hidden_5d=mix_speech_hidden_5d.expand(config.BATCH_SIZE,top_k_num,mix_speech_len,speech_fre,config.EMBEDDING_SIZE).contiguous()
mix_speech_hidden_5d_last=mix_speech_hidden_5d.view(-1,mix_speech_len,speech_fre,config.EMBEDDING_SIZE)
# att_speech_layer=ATTENTION(config.EMBEDDING_SIZE,'align').cuda()
att_speech_layer=ATTENTION(config.EMBEDDING_SIZE,'dot').cuda()
att_multi_speech=att_speech_layer(mix_speech_hidden_5d_last,mix_speech_multiEmbs.view(-1,config.EMBEDDING_SIZE))
# print att_multi_speech.size()
att_multi_speech=att_multi_speech.view(config.BATCH_SIZE,top_k_num,mix_speech_len,speech_fre) # bs,num_labels,len,fre这个东西
# print att_multi_speech.size()
multi_mask=att_multi_speech
# multi_mask=(att_multi_speech>0.5)
# multi_mask=Variable(torch.from_numpy(np.float32(multi_mask.data.cpu().numpy()))).cuda()
# top_k_mask_mixspeech_multi=top_k_mask_mixspeech.view(config.BATCH_SIZE,top_k_num,1,1).expand(config.BATCH_SIZE,top_k_num,mix_speech_len,speech_fre)
# multi_mask=multi_mask*Variable(top_k_mask_mixspeech_multi).cuda()
x_input_map=Variable(torch.from_numpy(train_data['mix_feas'])).cuda()
# print x_input_map.size()
x_input_map_multi=x_input_map.view(config.BATCH_SIZE,1,mix_speech_len,speech_fre).expand(config.BATCH_SIZE,top_k_num,mix_speech_len,speech_fre)
bss_eval_fromGenMap(multi_mask,x_input_map,top_k_mask_mixspeech,dict_idx2spk,train_data,batch_idx)
# SDR_SUM = np.append(SDR_SUM, bss_test.cal('batch_output/', 2))
# print 'SDR_SUM (len:{}) for epoch {} : {}'.format(SDR_SUM.shape,epoch_idx,SDR_SUM.mean())
# 1/0
SDR_SUM_total = np.append(SDR_SUM_total, bss_test.cal('batch_output/', 2))
print 'SDR_SUM (len:{}) for epoch {} : {}'.format(SDR_SUM_total.shape,epoch_idx,SDR_SUM_total.mean())