405 строки
17 KiB
Python
405 строки
17 KiB
Python
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT license.
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from block_zoo import *
|
|
import copy
|
|
import logging
|
|
from utils.exceptions import ConfigurationError, LayerUndefinedError, LayerConfigUndefinedError
|
|
from queue import Queue
|
|
from utils.common_utils import transform_tensors2params, transfer_to_gpu
|
|
|
|
from block_zoo.Embedding import *
|
|
|
|
EMBED_LAYER_NAME = 'Embedding'
|
|
EMBED_LAYER_ID = 'embedding'
|
|
|
|
|
|
def get_conf(layer_id, layer_name, input_layer_ids, all_layer_configs, model_input_ids, use_gpu,
|
|
conf_dict=None, shared_conf=None, succeed_embedding_flag=False, output_layer_flag=False,
|
|
target_num=None, fixed_lengths=None):
|
|
""" get layer configuration
|
|
|
|
Args
|
|
layer_id: layer identifier
|
|
layer_name: name of layer such as BiLSTM
|
|
input_layer_ids (list): the inputs of current layer
|
|
all_layer_configs (dict): records the conf class of each layer.
|
|
model_input_ids (set): the inputs of the model, e.g. ['query', 'passage']
|
|
use_gpu:
|
|
conf_dict:
|
|
shared_conf: if fixed_lengths is not None, the output_dim of shared_conf should be corrected!
|
|
flag:
|
|
output_layer_flag:
|
|
target_num: used for inference the dimension of output space if someone declare a dimension of -1
|
|
fixed_lengths
|
|
Returns:
|
|
configuration class coresponds to the layer
|
|
|
|
"""
|
|
if shared_conf:
|
|
conf = copy.deepcopy(shared_conf)
|
|
else:
|
|
try:
|
|
conf_dict['use_gpu'] = use_gpu
|
|
|
|
# for classification tasks, we usually add a Linear layer to project the output to dimension of number of classes. If we don't know the #classes, we can use '-1' instead and we would calculate the number of classes from the corpus.
|
|
if layer_name == 'Linear':
|
|
if isinstance(conf_dict['hidden_dim'], list) and conf_dict['hidden_dim'][-1] == -1:
|
|
assert output_layer_flag is True, "Only in the last layer, hidden_dim == -1 is allowed!"
|
|
assert target_num is not None, "Number of targets should be given!"
|
|
conf_dict['hidden_dim'][-1] = target_num
|
|
elif isinstance(conf_dict['hidden_dim'], int) and conf_dict['hidden_dim'] == -1:
|
|
assert output_layer_flag is True, "Only in the last layer, hidden_dim == -1 is allowed!"
|
|
assert target_num is not None, "Number of targets should be given!"
|
|
conf_dict['hidden_dim'] = target_num
|
|
|
|
conf = eval(layer_name + "Conf")(**conf_dict)
|
|
except NameError as e:
|
|
raise LayerConfigUndefinedError("\"%sConf\" has not been defined" % layer_name)
|
|
|
|
# verify the rank consistence of joint layers
|
|
if layer_name == EMBED_LAYER_NAME:
|
|
# the embedding layer
|
|
pass
|
|
else:
|
|
# make sure all the inputs to current layer exist
|
|
for input_layer_id in input_layer_ids:
|
|
if not (input_layer_id in all_layer_configs or input_layer_id in model_input_ids):
|
|
raise ConfigurationError("The input %s of layer %s does not exist. Please define it before "
|
|
"defining layer %s!" % (input_layer_id, layer_id, layer_id))
|
|
|
|
former_output_ranks = [all_layer_configs[input_layer_id].output_rank if input_layer_id in all_layer_configs else all_layer_configs[EMBED_LAYER_ID].output_rank for input_layer_id in input_layer_ids]
|
|
# inference input_dim
|
|
conf.input_dims = [all_layer_configs[input_layer_id].output_dim if input_layer_id in all_layer_configs else all_layer_configs[EMBED_LAYER_ID].output_dim for input_layer_id in input_layer_ids]
|
|
|
|
# If the inputs come from embedding layer and fixed_lengths exist, set the length to input_dims
|
|
if len(input_layer_ids) == 1 and input_layer_ids[0] in model_input_ids and fixed_lengths:
|
|
conf.input_dims[0][1] = fixed_lengths[input_layer_ids[0]]
|
|
|
|
# check and verify input ranks
|
|
if conf.num_of_inputs > 0:
|
|
if conf.num_of_inputs != len(input_layer_ids):
|
|
raise ConfigurationError("%s only accept %d inputs but you feed %d inputs to it!" % \
|
|
(layer_name, conf.num_of_inputs, len(input_layer_ids)))
|
|
elif conf.num_of_inputs == -1:
|
|
conf.num_of_inputs = len(input_layer_ids)
|
|
if isinstance(conf.input_ranks, list):
|
|
conf.input_ranks = conf.input_ranks * conf.num_of_inputs
|
|
else:
|
|
logging.warning("[For developer of %s] The input_ranks attribute should be a list!" % (layer_name))
|
|
[conf.input_ranks] * conf.num_of_inputs
|
|
|
|
for input_rank, former_output_rank in zip(conf.input_ranks, former_output_ranks):
|
|
if input_rank != -1 and input_rank != former_output_rank:
|
|
raise ConfigurationError("Input ranks of %s are inconsistent with former layers" % layer_id)
|
|
conf.input_ranks = copy.deepcopy(former_output_ranks)
|
|
|
|
# inference and varification inside the layer
|
|
conf.inference() # update some attributes which relies on input dimension or something else
|
|
conf.verify() # verify if the configuration is legal
|
|
|
|
logging.debug('Layer id: %s; name: %s; input_dims: %s; input_ranks: %s; output_dim: %s; output_rank: %s' % (layer_id, layer_name, conf.input_dims if layer_id != 'embedding' else 'None', conf.input_ranks, conf.output_dim, conf.output_rank))
|
|
|
|
return conf
|
|
|
|
|
|
def get_layer(layer_name, conf):
|
|
"""
|
|
|
|
Args:
|
|
layer_name:
|
|
conf: configuration class
|
|
|
|
Returns:
|
|
specific layer
|
|
|
|
"""
|
|
try:
|
|
layer = eval(layer_name)(conf)
|
|
except NameError as e:
|
|
raise Exception("%s; Layer \"%s\" has not been defined" % (str(e), layer_name))
|
|
return layer
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, conf, problem, vocab_info, use_gpu):
|
|
"""
|
|
|
|
Args:
|
|
inputs: ['string1', 'string2']
|
|
layer_archs: The layers must produce tensors with similar shapes. The layers may be nested.
|
|
[
|
|
{
|
|
'layer': Layer name,
|
|
'conf': {xxxx}
|
|
},
|
|
[
|
|
{
|
|
'layer': Layer name,
|
|
'conf': {},
|
|
},
|
|
{
|
|
'layer': Layer name,
|
|
'conf': {},
|
|
}
|
|
]
|
|
]
|
|
vocab_info:
|
|
{
|
|
'word': {
|
|
'vocab_size': xxx,
|
|
'init_weights': np matrix
|
|
}
|
|
'postag': {
|
|
'vocab_size': xxx,
|
|
'init_weights': None
|
|
}
|
|
}
|
|
"""
|
|
super(Model, self).__init__()
|
|
|
|
inputs = conf.object_inputs_names
|
|
layer_archs = conf.architecture
|
|
target_num = problem.output_target_num()
|
|
|
|
# correct the real fixed length if begin/end of sentence are added
|
|
if conf.fixed_lengths:
|
|
fixed_lengths_corrected = copy.deepcopy(conf.fixed_lengths)
|
|
for seq in fixed_lengths_corrected:
|
|
if problem.with_bos_eos:
|
|
fixed_lengths_corrected[seq] += 2
|
|
else:
|
|
fixed_lengths_corrected = None
|
|
|
|
self.use_gpu = use_gpu
|
|
|
|
all_layer_configs = dict()
|
|
self.layers = nn.ModuleDict()
|
|
self.layer_inputs = dict()
|
|
self.layer_dependencies = dict()
|
|
self.layer_dependencies[EMBED_LAYER_ID] = set()
|
|
# change output_layer_id to list for support multi_output
|
|
self.output_layer_id = []
|
|
|
|
for layer_index, layer_arch in enumerate(layer_archs):
|
|
output_layer_flag = True if 'output_layer_flag' in layer_arch and layer_arch['output_layer_flag'] is True else False
|
|
succeed_embedding_flag = True if layer_index > 0 and 'inputs' in layer_arch and \
|
|
[input in inputs for input in layer_arch['inputs']].count(True) == len(layer_arch['inputs']) else False
|
|
|
|
if output_layer_flag:
|
|
self.output_layer_id.append(layer_arch['layer_id'])
|
|
# if hasattr(self, 'output_layer_id'):
|
|
# raise ConfigurationError("There should be only one output!")
|
|
# else:
|
|
# self.output_layer_id = layer_arch['layer_id']
|
|
|
|
if layer_index == 0:
|
|
# embedding layer
|
|
emb_conf = copy.deepcopy(vocab_info)
|
|
for input_cluster in emb_conf:
|
|
emb_conf[input_cluster]['dim'] = layer_arch['conf'][input_cluster]['dim']
|
|
emb_conf[input_cluster]['fix_weight'] = layer_arch['conf'][input_cluster].get('fix_weight', False)
|
|
|
|
all_layer_configs[EMBED_LAYER_ID] = get_conf(EMBED_LAYER_ID, layer_arch['layer'],
|
|
None, all_layer_configs, inputs, self.use_gpu, conf_dict={'conf': emb_conf},
|
|
shared_conf=None, succeed_embedding_flag=False, output_layer_flag=output_layer_flag,
|
|
target_num=target_num, fixed_lengths=fixed_lengths_corrected)
|
|
self.add_layer(EMBED_LAYER_ID, get_layer(layer_arch['layer'], all_layer_configs[EMBED_LAYER_ID]))
|
|
else:
|
|
if layer_arch['layer'] in self.layers and not 'conf' in layer_arch:
|
|
# reuse formly defined layers (share the same parameters)
|
|
logging.debug("Layer id: %s; Sharing configuration with layer %s" % (layer_arch['layer_id'], layer_arch['layer']))
|
|
conf_dict = None
|
|
shared_conf = all_layer_configs[layer_arch['layer']]
|
|
else:
|
|
conf_dict = layer_arch['conf']
|
|
shared_conf = None
|
|
|
|
# if the layer is EncoderDecoder, inference the vocab size
|
|
if layer_arch['layer'] == 'EncoderDecoder':
|
|
layer_arch['conf']['decoder_conf']['decoder_vocab_size'] = target_num
|
|
all_layer_configs[layer_arch['layer_id']] = get_conf(layer_arch['layer_id'], layer_arch['layer'],
|
|
layer_arch['inputs'], all_layer_configs, inputs, self.use_gpu, conf_dict=conf_dict,
|
|
shared_conf=shared_conf, succeed_embedding_flag=succeed_embedding_flag,
|
|
output_layer_flag=output_layer_flag, target_num=target_num,
|
|
fixed_lengths=fixed_lengths_corrected)
|
|
|
|
if layer_arch['layer'] in self.layers and not 'conf' in layer_arch:
|
|
self.add_layer(layer_arch['layer_id'], self.layers[layer_arch['layer']])
|
|
else:
|
|
self.add_layer(layer_arch['layer_id'], get_layer(layer_arch['layer'], all_layer_configs[layer_arch['layer_id']]))
|
|
|
|
self.layer_inputs[layer_arch['layer_id']] = layer_arch['inputs']
|
|
|
|
# register dependencies, except embeddings
|
|
cur_layer_depend = set()
|
|
for layer_depend_id in layer_arch['inputs']:
|
|
if not layer_depend_id in inputs:
|
|
cur_layer_depend.add(layer_depend_id)
|
|
self.add_dependency(layer_arch['layer_id'], cur_layer_depend)
|
|
|
|
logging.debug("Layer dependencies: %s" % repr(self.layer_dependencies))
|
|
|
|
if not hasattr(self, 'output_layer_id'):
|
|
raise ConfigurationError("Please define an output layer")
|
|
|
|
self.layer_topological_sequence = self.get_topological_sequence()
|
|
|
|
def add_layer(self, layer_id, layer):
|
|
""" register a layer
|
|
|
|
Args:
|
|
layer_id:
|
|
layer:
|
|
|
|
Returns:
|
|
|
|
"""
|
|
if layer_id in self.layers:
|
|
raise ConfigurationError("The layer id %s is not unique!")
|
|
else:
|
|
self.layers[layer_id] = layer
|
|
|
|
def add_dependency(self, layer_id, depend_layer_id):
|
|
""" add the layers have to be proceed before layer_id
|
|
|
|
Args:
|
|
layer_id:
|
|
depend_layer_id:
|
|
|
|
Returns:
|
|
|
|
"""
|
|
if not layer_id in self.layer_dependencies:
|
|
self.layer_dependencies[layer_id] = set()
|
|
|
|
if isinstance(depend_layer_id, int):
|
|
self.layer_dependencies[layer_id].add(depend_layer_id)
|
|
else:
|
|
self.layer_dependencies[layer_id] |= set(depend_layer_id)
|
|
|
|
def remove_dependency(self, depend_layer_id):
|
|
""" remove dependencies on layer_id
|
|
|
|
Args:
|
|
layer_id:
|
|
|
|
Returns:
|
|
|
|
"""
|
|
for layer_id in self.layer_dependencies:
|
|
self.layer_dependencies[layer_id].remove(depend_layer_id)
|
|
|
|
def get_topological_sequence(self):
|
|
""" get topological sequence of nodes in the model
|
|
|
|
Returns:
|
|
|
|
"""
|
|
total_layer_ids = Queue()
|
|
for layer_id in self.layers.keys():
|
|
if layer_id != EMBED_LAYER_ID:
|
|
total_layer_ids.put(layer_id)
|
|
|
|
topological_list = []
|
|
circular_cnt = 0 # used for checking if there is at least one legal topological sorting
|
|
while not total_layer_ids.empty():
|
|
layer_id = total_layer_ids.get()
|
|
if len(self.layer_dependencies[layer_id]) == 0:
|
|
for layer_id2 in self.layer_dependencies:
|
|
if layer_id in self.layer_dependencies[layer_id2]:
|
|
self.layer_dependencies[layer_id2].remove(layer_id)
|
|
circular_cnt = 0
|
|
topological_list.append(layer_id)
|
|
else:
|
|
total_layer_ids.put(layer_id)
|
|
circular_cnt += 1
|
|
if circular_cnt >= total_layer_ids.qsize():
|
|
rest_layers = []
|
|
while not total_layer_ids.empty():
|
|
rest_layers.append(total_layer_ids.get())
|
|
raise ConfigurationError("The model architecture is illegal because there is a circular dependency "
|
|
"or there are some isolated layers. The layers can not be resolved: [%s]" % (", ".join(rest_layers)))
|
|
|
|
logging.debug("Topological sequence of nodes: %s" % (",".join(topological_list)))
|
|
return topological_list
|
|
|
|
def forward(self, inputs_desc, length_desc, *param_list):
|
|
"""
|
|
|
|
Args:
|
|
with the help of transform_tensors2params(inputs_desc, length_desc, param_list), we can get the below inputs and lengths
|
|
|
|
inputs: dict.
|
|
{
|
|
"string1":{
|
|
'word': word ids, [batch size, seq len]
|
|
'postag': postag ids,[batch size, seq len]
|
|
...
|
|
}
|
|
"string2":{
|
|
'word': word ids,[batch size, seq len]
|
|
'postag': postag ids,[batch size, seq len]
|
|
...
|
|
}
|
|
}
|
|
lengths: dict.
|
|
{
|
|
"string1": [...]
|
|
"string2": [...]
|
|
}
|
|
|
|
Returns:
|
|
|
|
"""
|
|
inputs, lengths = transform_tensors2params(inputs_desc, length_desc, param_list)
|
|
|
|
representation = dict()
|
|
representation[EMBED_LAYER_ID] = dict()
|
|
repre_lengths = dict()
|
|
repre_lengths[EMBED_LAYER_ID] = dict()
|
|
|
|
for input in inputs:
|
|
representation[input] = self.layers[EMBED_LAYER_ID](inputs[input], use_gpu=self.is_cuda())
|
|
if self.use_gpu:
|
|
repre_lengths[input] = transfer_to_gpu(lengths[input])
|
|
else:
|
|
repre_lengths[input] = lengths[input]
|
|
|
|
for layer_id in self.layer_topological_sequence:
|
|
#logging.debug("To proces layer %s" % layer_id)
|
|
input_params = []
|
|
for input_layer_id in self.layer_inputs[layer_id]:
|
|
input_params.append(representation[input_layer_id])
|
|
input_params.append(repre_lengths[input_layer_id])
|
|
|
|
representation[layer_id], repre_lengths[layer_id] = self.layers[layer_id](*input_params)
|
|
#logging.debug("Layer %s processed. output size: %s" % (layer_id, representation[layer_id].size()))
|
|
|
|
# for support multi_output
|
|
representation_output = dict()
|
|
for single_output_layer_id in self.output_layer_id:
|
|
representation_output[single_output_layer_id] = representation[single_output_layer_id]
|
|
return representation_output
|
|
|
|
def is_cuda(self):
|
|
return next(self.parameters()).data.is_cuda
|
|
|
|
def update_use_gpu(self, new_use_gpu):
|
|
self.use_gpu = new_use_gpu
|
|
for layer_id in self.layers.keys():
|
|
if isinstance(self.layers[layer_id], Embedding):
|
|
for input_cluster in self.layers[layer_id].embeddings:
|
|
if isinstance(self.layers[layer_id].embeddings[input_cluster], CNNCharEmbedding):
|
|
self.layers[layer_id].embeddings[input_cluster].layer_conf.use_gpu = new_use_gpu
|
|
elif isinstance(self.layers[layer_id], EncoderDecoder):
|
|
self.layers[layer_id].encoder.layer_conf.use_gpu = new_use_gpu
|
|
self.layers[layer_id].decoder.layer_conf.use_gpu = new_use_gpu
|
|
else:
|
|
self.layers[layer_id].layer_conf.use_gpu = new_use_gpu
|
|
|
|
|
|
|