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