fix gpu run

atulkum · atulkum · commit 34881ca37b1b · 2019-01-09T22:13:30.000-08:00
diff --git a/neural_ner/process_training.py b/neural_ner/process_training.py
@@ -155,7 +155,7 @@ def evaluate(self, data_type, num_samples=None):
 
 if __name__ == "__main__":
     from config import config
-    config.gpu = config.gpu and torch.cuda.is_available()
+    config.is_cuda = config.is_cuda and torch.cuda.is_available()
     mode = sys.argv[1]
     model_file_path = None
     if len(sys.argv) > 2:
diff --git a/neural_ner/train_utils.py b/neural_ner/train_utils.py
@@ -76,8 +76,8 @@ def batch_update(self, batch, pred):
         s_lengths = batch['words_lens']
         y = batch['tags']
         raw_sentence = batch['raw_sentence']
-        y_reals = y.data.numpy()
-        y_preds = pred.data.numpy()
+        y_reals = y.cpu().data.numpy()
+        y_preds = pred.cpu().data.numpy()
         for i, s_len in enumerate(s_lengths):
             r_tags = []
             p_tags = []
@@ -94,7 +94,7 @@ def batch_update(self, batch, pred):
                 self.predictions.append(new_line)
             self.predictions.append("")
 
-    def get_metric(self, log_dir):
+    def get_metric(self, log_dir, is_cf=False):
         # Write predictions to disk and run CoNLL script externally
         eval_id = np.random.randint(1000000, 2000000)
         output_path = os.path.join(log_dir, "eval.%i.output" % eval_id)
@@ -106,25 +106,26 @@ def get_metric(self, log_dir):
 
         # CoNLL evaluation results
         eval_lines = [l.rstrip() for l in codecs.open(scores_path, 'r', 'utf8')]
-        for line in eval_lines:
-            print line
-
-        # Confusion matrix with accuracy for each tag
-        print ("{: >2}{: >7}{: >7}%s{: >9}" % ("{: >7}" * self.n_tags)).format(
-            "ID", "NE", "Total",
-            *([self.vocab.id_to_tag[i] for i in xrange(self.n_tags)] + ["Percent"])
-        )
-        for i in xrange(self.n_tags):
+        if is_cf:
+            for line in eval_lines:
+                print line
+
+            # Confusion matrix with accuracy for each tag
             print ("{: >2}{: >7}{: >7}%s{: >9}" % ("{: >7}" * self.n_tags)).format(
-                str(i), self.vocab.id_to_tag[i], str(self.count[i].sum()),
-                *([self.count[i][j] for j in xrange(self.n_tags)] +
-                  ["%.3f" % (self.count[i][i] * 100. / max(1, self.count[i].sum()))])
+                "ID", "NE", "Total",
+                *([self.vocab.id_to_tag[i] for i in xrange(self.n_tags)] + ["Percent"])
+            )
+            for i in xrange(self.n_tags):
+                print ("{: >2}{: >7}{: >7}%s{: >9}" % ("{: >7}" * self.n_tags)).format(
+                    str(i), self.vocab.id_to_tag[i], str(self.count[i].sum()),
+                    *([self.count[i][j] for j in xrange(self.n_tags)] +
+                      ["%.3f" % (self.count[i][i] * 100. / max(1, self.count[i].sum()))])
             )
 
-        # Global accuracy
-        print "%i/%i (%.5f%%)" % (
-            self.count.trace(), self.count.sum(), 100. * self.count.trace() / max(1, self.count.sum())
-        )
+            # Global accuracy
+            print "%i/%i (%.5f%%)" % (
+                self.count.trace(), self.count.sum(), 100. * self.count.trace() / max(1, self.count.sum())
+            )
 
         # F1 on all entities
         return float(eval_lines[1].strip().split()[-1])
