print('unique labels = {}'.format(dataset.get_unique_labels_num()))
print('vocab size = {}'.format(dataset.get_vocab_size()))
if config.model_name == 'CNN':
model = Classification(config.embedding, config.strmaxlen,
dataset.get_unique_labels_num(),
dataset.get_vocab_size())
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)
elif config.model_name == 'RCNN':
model = RCNN(config.embedding, config.strmaxlen,
dataset.get_unique_labels_num(), dataset.get_vocab_size())
if config.mode == 'train':
model = model.cuda()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)
elif config.model_name == "DUALRCNN":
model = DualRCNN(config.embedding, config.strmaxlen,
dataset.get_unique_labels_num(),
dataset.get_vocab_size())
if torch.cuda.is_available():
model = model.cuda()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)
elif config.model_name == 'TRANSFORMER':
model = Model(config2)
