refactor

atulkum · atulkum · commit 7d4b17ec130c · 2019-01-19T21:23:00.000-08:00
diff --git a/neural_ner/config.py b/neural_ner/config.py
@@ -1,5 +1,11 @@
 import os
 import torch
+import numpy as np
+
+np.random.seed(123)
+torch.manual_seed(123)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(123)
 
 class Config(object):
     pass
@@ -18,7 +24,7 @@ class Config(object):
 config.dropout_ratio = 0.5
 
 config.max_grad_norm = 5.0
-config.batch_size = 32
+config.batch_size = 1
 config.num_epochs = 100
 
 config.print_every = 100
@@ -45,4 +51,9 @@ class Config(object):
 
 config.is_l2_loss = False
 
-config.is_cuda = config.is_cuda and torch.cuda.is_available()
+config.model_name = 'model.NER_SOFTMAX_CHAR'
+config.optimizer = 'sgd'
+
+# config postprocess
+config.is_cuda = config.is_cuda and torch.cuda.is_available()
+
diff --git a/neural_ner/model.py b/neural_ner/model.py
@@ -15,26 +15,6 @@
 
 logging.basicConfig(level=logging.INFO)
 
-np.random.seed(123)
-torch.manual_seed(123)
-if torch.cuda.is_available():
-    torch.cuda.manual_seed_all(123)
-
-def get_model(vocab, config, model_file_path, is_eval=False):
-    #model = NER_SOFTMAX_CHAR_CRF(vocab, config)
-    model = NER_SOFTMAX_CHAR(vocab, config)
-
-    if is_eval:
-        model = model.eval()
-    if config.is_cuda:
-        model = model.cuda()
-
-    if model_file_path is not None:
-        state = torch.load(model_file_path, map_location=lambda storage, location: storage)
-        model.load_state_dict(state['model'], strict=False)
-
-    return model
-
 class NER_SOFTMAX_CHAR(nn.Module):
     def __init__(self, vocab, config):
         super(NER_SOFTMAX_CHAR, self).__init__()
diff --git a/neural_ner/model_utils.py b/neural_ner/model_utils.py
@@ -1,7 +1,44 @@
-from __future__ import unicode_literals, print_function, division
-
 import torch
 import numpy as np
+import importlib
+import logging
+from torch.optim import Adam, SGD
+
+logging.basicConfig(level=logging.INFO)
+
+def get_model(vocab, config, model_file_path, is_eval=False):
+    class_path = config.model_name.split('.')
+    class_name = class_path[-1]
+    class_module = '.'.join(class_path[:-1])
+
+    ModelClass = getattr(importlib.import_module(class_module), class_name)
+
+    model = ModelClass(vocab, config)
+
+    if is_eval:
+        model = model.eval()
+    if config.is_cuda:
+        model = model.cuda()
+
+    if model_file_path is not None:
+        state = torch.load(model_file_path, map_location=lambda storage, location: storage)
+        model.load_state_dict(state['model'], strict=False)
+
+    return model
+
+def get_optimizer(model, config):
+    params = list(filter(lambda p: p.requires_grad, model.parameters()))
+
+    optimizer = None
+    if config.optimizer == 'adam':
+        optimizer = Adam(params, amsgrad=True)
+    elif config.optimizer == 'sgd':
+        optimizer = SGD(params, lr=0.01, momentum=0.9)
+
+    num_params = sum(p.numel() for p in params)
+    logging.info("Number of params: %d" % num_params)
+
+    return optimizer, params
 
 def get_mask(lengths, config):
     seq_lens = lengths.view(-1, 1)
diff --git a/neural_ner/process_training.py b/neural_ner/process_training.py
@@ -5,11 +5,10 @@
 
 import torch
 from torch.nn.utils import clip_grad_norm_
-from torch.optim import Adam, SGD
 
 from data_utils.batcher import DatasetConll2003
 from data_utils.vocab import Vocab
-from model import get_model
+from model_utils import get_model, get_optimizer
 from train_utils import setup_train_dir, save_model, write_summary, \
     get_param_norm, get_grad_norm, Evaluter
 
@@ -49,12 +48,7 @@ def train_one_batch(self, batch, optimizer, params):
 
     def train(self):
         train_dir, summary_writer = setup_train_dir(self.config)
-
-        params = list(filter(lambda p: p.requires_grad, self.model.parameters()))
-        optimizer = Adam(params, amsgrad=True)
-
-        num_params = sum(p.numel() for p in params)
-        logging.info("Number of params: %d" % num_params)
+        optimizer, params = get_optimizer(self.model, self.config)
 
         exp_loss, best_dev_f1 = None, None
 
