Merge pull request Shawn1993#12 from oneTaken/master

update details for better use

Author: Shawn1993

main.py

@@ -19,8 +19,10 @@
 parser.add_argument('-test-interval', type=int, default=100, help='how many steps to wait before testing [default: 100]')
 parser.add_argument('-save-interval', type=int, default=500, help='how many steps to wait before saving [default:500]')
 parser.add_argument('-save-dir', type=str, default='snapshot', help='where to save the snapshot')
+parser.add_argument('-early-stop', type=int, default=1000, help='iteration numbers to stop without performance increasing')
+parser.add_argument('-save-best', type=bool, default=True, help='whether to save when get best performance')
 # data 
-parser.add_argument('-shuffle', action='store_true', default=False, help='shuffle the data every epoch' )
+parser.add_argument('-shuffle', action='store_true', default=False, help='shuffle the data every epoch')
 # model
 parser.add_argument('-dropout', type=float, default=0.5, help='the probability for dropout [default: 0.5]')
 parser.add_argument('-max-norm', type=float, default=3.0, help='l2 constraint of parameters [default: 3.0]')
@@ -30,7 +32,7 @@
 parser.add_argument('-static', action='store_true', default=False, help='fix the embedding')
 # device
 parser.add_argument('-device', type=int, default=-1, help='device to use for iterate data, -1 mean cpu [default: -1]')
-parser.add_argument('-no-cuda', action='store_true', default=False, help='disable the gpu' )
+parser.add_argument('-no-cuda', action='store_true', default=False, help='disable the gpu')
 # option
 parser.add_argument('-snapshot', type=str, default=None, help='filename of model snapshot [default: None]')
 parser.add_argument('-predict', type=str, default=None, help='predict the sentence given')
@@ -69,7 +71,7 @@ def mr(text_field, label_field, **kargs):
 text_field = data.Field(lower=True)
 label_field = data.Field(sequential=False)
 train_iter, dev_iter = mr(text_field, label_field, device=-1, repeat=False)
-#train_iter, dev_iter, test_iter = sst(text_field, label_field, device=-1, repeat=False)
+# train_iter, dev_iter, test_iter = sst(text_field, label_field, device=-1, repeat=False)
 
 
 # update args and print
@@ -85,33 +87,30 @@ def mr(text_field, label_field, **kargs):
 
 
 # model
-if args.snapshot is None:
-    cnn = model.CNN_Text(args)
-else :
-    print('\nLoading model from [%s]...' % args.snapshot)
-    try:
-        cnn = torch.load(args.snapshot)
-    except :
-        print("Sorry, This snapshot doesn't exist."); exit()
+cnn = model.CNN_Text(args)
+if args.snapshot is not None:
+    print('\nLoading model from {}...'.format(args.snapshot))
+    cnn.load_state_dict(torch.load(args.snapshot))
 
 if args.cuda:
+    torch.cuda.set_device(args.device)
     cnn = cnn.cuda()
 
 
 # train or predict
 if args.predict is not None:
     label = train.predict(args.predict, cnn, text_field, label_field, args.cuda)
     print('\n[Text]  {}\n[Label] {}\n'.format(args.predict, label))
-elif args.test :
+elif args.test:
     try:
         train.eval(test_iter, cnn, args) 
     except Exception as e:
         print("\nSorry. The test dataset doesn't  exist.\n")
-else :
+else:
     print()
     try:
         train.train(train_iter, dev_iter, cnn, args)
     except KeyboardInterrupt:
-        print('-' * 89)
+        print('\n' + '-' * 89)
         print('Exiting from training early')
 

model.py

@@ -1,11 +1,13 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch.autograd import Variable
 
-class  CNN_Text(nn.Module):
+
+class CNN_Text(nn.Module):
 
     def __init__(self, args):
-        super(CNN_Text,self).__init__()
+        super(CNN_Text, self).__init__()
         self.args = args
 
         V = args.embed_num
@@ -16,7 +18,7 @@ def __init__(self, args):
         Ks = args.kernel_sizes
 
         self.embed = nn.Embedding(V, D)
-        #self.convs1 = [nn.Conv2d(Ci, Co, (K, D)) for K in Ks]
+        # self.convs1 = [nn.Conv2d(Ci, Co, (K, D)) for K in Ks]
         self.convs1 = nn.ModuleList([nn.Conv2d(Ci, Co, (K, D)) for K in Ks])
         '''
         self.conv13 = nn.Conv2d(Ci, Co, (3, D))
@@ -27,23 +29,21 @@ def __init__(self, args):
         self.fc1 = nn.Linear(len(Ks)*Co, C)
 
     def conv_and_pool(self, x, conv):
-        x = F.relu(conv(x)).squeeze(3) #(N,Co,W)
+        x = F.relu(conv(x)).squeeze(3)  # (N, Co, W)
         x = F.max_pool1d(x, x.size(2)).squeeze(2)
         return x
 
-
     def forward(self, x):
-        x = self.embed(x) # (N,W,D)
+        x = self.embed(x)  # (N, W, D)
 
         if self.args.static:
             x = Variable(x)
 
-        x = x.unsqueeze(1) # (N,Ci,W,D)
-
-        x = [F.relu(conv(x)).squeeze(3) for conv in self.convs1] #[(N,Co,W), ...]*len(Ks)
+        x = x.unsqueeze(1)  # (N, Ci, W, D)
 
+        x = [F.relu(conv(x)).squeeze(3) for conv in self.convs1]  # [(N, Co, W), ...]*len(Ks)
 
-        x = [F.max_pool1d(i, i.size(2)).squeeze(2) for i in x] #[(N,Co), ...]*len(Ks)
+        x = [F.max_pool1d(i, i.size(2)).squeeze(2) for i in x]  # [(N, Co), ...]*len(Ks)
 
         x = torch.cat(x, 1)
 
@@ -53,6 +53,6 @@ def forward(self, x):
         x3 = self.conv_and_pool(x,self.conv15) #(N,Co)
         x = torch.cat((x1, x2, x3), 1) # (N,len(Ks)*Co)
         '''
-        x = self.dropout(x) # (N,len(Ks)*Co)
-        logit = self.fc1(x) # (N,C)
+        x = self.dropout(x)  # (N, len(Ks)*Co)
+        logit = self.fc1(x)  # (N, C)
         return logit

mydatasets.py

@@ -2,7 +2,7 @@
 import os
 import random
 import tarfile
-from six.moves import urllib
+import urllib
 from torchtext import data
 
 
@@ -86,7 +86,7 @@ def clean_str(string):
         super(MR, self).__init__(examples, fields, **kwargs)
 
     @classmethod
-    def splits(cls, text_field, label_field, dev_ratio=.1, shuffle=True ,root='.', **kwargs):
+    def splits(cls, text_field, label_field, dev_ratio=.1, shuffle=True, root='.', **kwargs):
         """Create dataset objects for splits of the MR dataset.
 
         Arguments:

train.py

@@ -12,6 +12,8 @@ def train(train_iter, dev_iter, model, args):
     optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
 
     steps = 0
+    best_acc = 0
+    last_step = 0
     model.train()
     for epoch in range(1, args.epochs+1):
         for batch in train_iter:
@@ -40,12 +42,17 @@ def train(train_iter, dev_iter, model, args):
                                                                              corrects,
                                                                              batch.batch_size))
             if steps % args.test_interval == 0:
-                eval(dev_iter, model, args)
-            if steps % args.save_interval == 0:
-                if not os.path.isdir(args.save_dir): os.makedirs(args.save_dir)
-                save_prefix = os.path.join(args.save_dir, 'snapshot')
-                save_path = '{}_steps{}.pt'.format(save_prefix, steps)
-                torch.save(model, save_path)
+                dev_acc = eval(dev_iter, model, args)
+                if dev_acc > best_acc:
+                    best_acc = dev_acc
+                    last_step = steps
+                    if args.save_best:
+                        save(model, args.save_dir, 'best', steps)
+                else:
+                    if steps - last_step >= args.early_stop:
+                        print('early stop by {} steps.'.format(args.early_stop))
+            elif steps % args.save_interval == 0:
+                save(model, args.save_dir, 'snapshot', steps)
 
 
 def eval(data_iter, model, args):
@@ -65,13 +72,13 @@ def eval(data_iter, model, args):
                      [1].view(target.size()).data == target.data).sum()
 
     size = len(data_iter.dataset)
-    avg_loss = avg_loss/size
+    avg_loss /= size
     accuracy = 100.0 * corrects/size
-    model.train()
     print('\nEvaluation - loss: {:.6f}  acc: {:.4f}%({}/{}) \n'.format(avg_loss, 
                                                                        accuracy, 
                                                                        corrects, 
                                                                        size))
+    return accuracy
 
 
 def predict(text, model, text_field, label_feild, cuda_flag):
@@ -83,8 +90,16 @@ def predict(text, model, text_field, label_feild, cuda_flag):
     x = text_field.tensor_type(text)
     x = autograd.Variable(x, volatile=True)
     if cuda_flag:
-        x =x.cuda()
+        x = x.cuda()
     print(x)
     output = model(x)
     _, predicted = torch.max(output, 1)
     return label_feild.vocab.itos[predicted.data[0][0]+1]
+
+
+def save(model, save_dir, save_prefix, steps):
+    if not os.path.isdir(save_dir):
+        os.makedirs(save_dir)
+    save_prefix = os.path.join(save_dir, save_prefix)
+    save_path = '{}_steps_{}.pt'.format(save_prefix, steps)
+    torch.save(model.state_dict(), save_path)