def model_selector(config, model_id, use_element):
model = None
if use_element:
print("use element")
model = ModelWithElement(config, model_id)
elif model_id == 0:
model = FastText(config)
elif model_id == 1:
model = TextCNN(config)
elif model_id == 2:
model = TextRCNN(config)
elif model_id == 3:
model = TextRNN(config)
elif model_id == 4:
model = HAN(config)
elif model_id == 5:
model = CNNWithDoc2Vec(config)
elif model_id == 6:
model = RCNNWithDoc2Vec(config)
elif model_id == 7:
model = CNNwithInception(config)
else:
print("Input ERROR!")
exit(2)
return model
def load_model(model_path, model_id, config):
if model_id == 0:
model = FastText(config)
elif model_id == 1:
model = TextCNN(config)
elif model_id == 2:
model = TextRCNN(config)
elif model_id == 4:
model = HAN(config)
# print(model)
# time_stamp = '1510844987'
# final_model_path = config.model_path+"."+time_stamp+"."+config.model_names[model_id]
print("load model data:", model_path)
model.load_state_dict(torch.load(model_path))
if config.has_cuda:
model = model.cuda()
return model
def _model_selector(config, model_id):
model = None
if model_id == 0:
model = FastText(config)
elif model_id == 1:
model = TextCNN(config)
elif model_id == 2:
model = TextRCNN(config)
elif model_id == 4:
model = HAN(config)
elif model_id == 5:
model = CNNWithDoc2Vec(config)
elif model_id == 6:
model = RCNNWithDoc2Vec(config)
else:
print("Input ERROR!")
exit(2)
return model
def get_models(
vocab_size, # 词典大小
n_class=10, # 类别个数
seq_len=38, # 句子长度
device=None): # 设备
""" 获取所有需要训练的模型 """
if device is None:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
fast_text = FastText(vocab_size=vocab_size, n_class=n_class)
text_cnn = TextCNN(vocab_size=vocab_size, n_class=n_class)
text_rnn = TextRNN(vocab_size=vocab_size, n_class=n_class)
text_rcnn = TextRCNN(vocab_size=vocab_size, n_class=n_class)
transformer = Transformer(vocab_size=vocab_size,
seq_len=seq_len,
n_class=n_class,
device=device)
return [fast_text, text_cnn, text_rnn, text_rcnn, transformer]
def __call__(self, num_labels: int):
if self.choose_pretrain == "Bert":
if self.resume_path:
model_dir = self.resume_path
else:
model_dir = config.bert_model_dir
if self.choose_model == "BertFC":
model = BertFCForMultiLable.from_pretrained(
model_dir, num_labels=num_labels)
elif self.choose_model == "BertCNN":
model = BertCNNForMultiLabel.from_pretrained(
model_dir, num_labels=num_labels)
elif self.choose_model == "BertRCNN":
model = BertRCNNForMultiLabel.from_pretrained(
model_dir, num_labels=num_labels)
elif self.choose_model == "BertDPCNN":
model = BertDPCNNForMultiLabel.from_pretrained(
model_dir, num_labels=num_labels)
else:
raise ModelNotDefinedError
elif self.choose_pretrain in ["Word2vec", "Nopretrain"]:
if self.resume_path:
model_dir = self.resume_path
else:
model_dir = None
if self.choose_pretrain == "Word2vec":
pretrain_model_dir = config.word2vec_model_dir
else:
pretrain_model_dir = None
if self.choose_model == "TextCNN":
cnn_config = config.cnn
cnn_config.embedding_pretrained = pretrain_model_dir
cnn_config.embedding_size = config.embedding_size
cnn_config.vocab_size = config.vocab_size
cnn_config.dropout = config.dropout
cnn_config.num_labels = num_labels
if self.resume_path:
model = TextCNN(cnn_config)
state_dict_file = os.path.join(model_dir,
"pytorch_model.bin")
model.load_state_dict(torch.load(state_dict_file))
else:
model = TextCNN(cnn_config)
elif self.choose_model == "TextRCNN":
rcnn_config = config.rcnn
rcnn_config.num_labels = num_labels
if self.resume_path:
model = TextRCNN(rcnn_config)
else:
model = TextRCNN(rcnn_config)
model.load_state_dict(torch.load(model_dir))
else:
raise ModelNotDefinedError
else:
raise ModelNotDefinedError
return model