Added the basic training code for attention models

Author: digbose92

Visual_Attention/attention_models.py

@@ -22,10 +22,7 @@ def forward(self, v, q):
     def logits(self, v, q):
         num_objs = v.size(1)
         q = q.unsqueeze(1).repeat(1, num_objs, 1)
-        print(q.size())
-        print(v.size())
         vq = torch.cat((v, q), 2)
-        print(vq.size())
         joint_repr = self.nonlinear(vq)
         logits = self.linear(joint_repr)
         return logits

Visual_Attention/train_models.py

@@ -23,6 +23,7 @@
 from vqa_dataset_attention import *
 import torch.nn as nn
 import random
+import utils
 
 def instance_bce_with_logits(logits, labels):
     assert logits.dim() == 2
@@ -38,6 +39,56 @@ def compute_score_with_logits(logits, labels):
     scores = (one_hots * labels)
     return scores
 
+def evaluate_model(model, valid_dataloader,device):
+    score = 0
+    Validation_loss = 0
+    upper_bound = 0
+    num_data = 0
+    V_loss=0 
+    print('Validation started')
+    #i, (feat, quest, label, target) 
+    for data in tqdm(valid_dataloader):
+
+        feat, quest, label, target = data
+        feat = feat.to(device)
+        quest = quest.to(device)
+        target = target.to(device) # true labels
+
+        pred = model(feat, quest, target)
+        loss = instance_bce_with_logits(pred, target)
+        V_loss += loss.item() * feat.size(0)
+        batch_score = compute_score_with_logits(pred, target.data).sum()
+        score += batch_score
+        upper_bound += (target.max(1)[0]).sum()
+        num_data += pred.size(0)
+        
+    score = score / len(valid_dataloader.dataset)
+    V_loss /= len(valid_dataloader.dataset)
+    upper_bound = upper_bound / len(valid_dataloader.dataset)
+    print(score,V_loss)
+    return score, upper_bound, V_loss
+
+def single_batch_run(model,train_dataloader,valid_dataloader,device,output_folder,optim):
+    feat_train, quest_train, label_train, target_train = next(iter(train_dataloader))
+    feat_train = feat_train.to(device_select)
+    quest_train = quest_train.to(device_select)
+    target_train = target_train.to(device_select) # true labels
+    pred = model(feat_train, quest_train, target_train)
+    loss = instance_bce_with_logits(pred, target_train)
+    logger = utils.Logger(os.path.join(output_folder, 'log_single_batch.txt'))
+    #print(loss)
+    loss.backward()
+    nn.utils.clip_grad_norm_(model.parameters(), 0.25)
+    optim.step()
+    optim.zero_grad()
+    batch_score = compute_score_with_logits(pred, target_train.data).sum()
+    model.train(False)
+    eval_score, bound, V_loss = evaluate_model(model, valid_dataloader,device)
+    model.train(True)
+    #logger.write('epoch %d, time: %.2f' % (epoch, time.time()-t))
+    #logger.write('\ttrain_loss: %.3f, score: %.3f' % (total_loss, train_score))
+    logger.write('\teval loss: %.3f, score: %.3f (%.3f)' % (V_loss, 100 * eval_score, 100 * bound))
+
 def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument('--eval', action='store_true', help='set this to evaluate.')
@@ -52,7 +103,7 @@ def parse_args():
     parser.add_argument('--norm', type=str, default='weight', help='weight, batch, layer, none')
     parser.add_argument('--model', type=str, default='A3x2')
     parser.add_argument('--output', type=str, default='saved_models/')
-    parser.add_argument('--batch_size', type=int, default=128)
+    parser.add_argument('--batch_size', type=int, default=512)
     parser.add_argument('--weight_decay', type=float, default=0)
     parser.add_argument('--optimizer', type=str, default='Adamax', help='Adam, Adamax, Adadelta, RMSprop')
     parser.add_argument('--initializer', type=str, default='kaiming_normal')
@@ -67,11 +118,14 @@ def parse_args():
     feats_data_path="/data/digbose92/VQA/COCO/train_hdf5_COCO/"
     data_root="/proj/digbose92/VQA/VisualQuestion_VQA/common_resources"
     npy_file="../../VisualQuestion_VQA/Visual_All/data/glove6b_init_300d.npy"
+    output_folder="/proj/digbose92/VQA/VisualQuestion_VQA/Visual_Attention/results"
     seed = 0
     args = parse_args()
     #device_selection
-    device=1
-    torch.cuda.set_device(device)
+    device_ids=[0,1]
+    #device_select=1
+    #torch.cuda.set_device(device_select)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     if args.seed == 0:
         seed = random.randint(1, 10000)
@@ -84,16 +138,22 @@ def parse_args():
         torch.cuda.manual_seed(args.seed)
 
     #train dataset
-    train_dataset=Dataset_VQA(img_root_dir=image_root_dir,feats_data_path=feats_data_path,dictionary=dictionary,dataroot=data_root,arch_choice="resnet152",layer_option="pool")
-    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=8)
+    train_dataset=Dataset_VQA(img_root_dir=image_root_dir,feats_data_path=feats_data_path,dictionary=dictionary,choice='train',dataroot=data_root,arch_choice="resnet152",layer_option="pool")
+    valid_dataset=Dataset_VQA(img_root_dir=image_root_dir,feats_data_path=feats_data_path,dictionary=dictionary,choice='val',dataroot=data_root,arch_choice="resnet152",layer_option="pool")
+    
+    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=10)
+    val_loader=DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False, num_workers=8)
+    print(len(train_loader))
+    print(len(val_loader))
     total_step=len(train_loader)
 
     #model related issues
     model = attention_baseline(train_dataset, num_hid=args.num_hid, dropout= args.dropout, norm=args.norm,\
                                activation=args.activation, drop_L=args.dropout_L, drop_G=args.dropout_G,\
                                drop_W=args.dropout_W, drop_C=args.dropout_C)
 
-    model=model.to(device)
+    #model=model.to(device_select)
+    
 
     if args.initializer == 'xavier_normal':
         model.apply(weights_init_xn)
@@ -105,7 +165,9 @@ def parse_args():
         model.apply(weights_init_ku)
 
     model.w_emb.init_embedding(npy_file)
-
+    if torch.cuda.device_count() > 1:
+        print("Let's use", torch.cuda.device_count(), "GPUs!")
+        model=torch.nn.DataParallel(model, device_ids=device_ids).to(device)
 
     if args.optimizer == 'Adadelta':
         optim = torch.optim.Adadelta(model.parameters(), rho=0.95, eps=1e-6, weight_decay=args.weight_decay)
@@ -115,39 +177,70 @@ def parse_args():
         optim = torch.optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.weight_decay)
     else:
         optim = torch.optim.Adamax(model.parameters(), weight_decay=args.weight_decay)
-
+    
+    logger = utils.Logger(os.path.join(output_folder, 'log.txt'))
+    best_eval_score = 0
     print('Starting training')
+    
+    #placeholder for checking training and testuing working or not
+    #single_batch_run(model,train_loader,val_loader,device_select,output_folder,optim)
+
+    device_select=0
+
     for epoch in range(args.epochs):
         total_loss = 0
         train_score = 0
         t = time.time()
         correct = 0
         step=0
+        start_time=time.time()
         for i, (feat, quest, label, target) in enumerate(train_loader):
+
             feat = feat.to(device)
             quest = quest.to(device)
             target = target.to(device) # true labels
 
             pred = model(feat, quest, target)
             loss = instance_bce_with_logits(pred, target)
-            print(loss)
+            #print(loss)
             loss.backward()
-            nn.utils.clip_grad_norm(model.parameters(), 0.25)
+            nn.utils.clip_grad_norm_(model.parameters(), 0.25)
             optim.step()
             optim.zero_grad()
 
             batch_score = compute_score_with_logits(pred, target.data).sum()
             total_loss += loss.item() * feat.size(0)
             train_score += batch_score
             if(step%10==0):
-            #optimizer.zero_grad()
-                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
-                    .format(epoch, args.epochs, step, total_step, loss.item()))
+                end_time=time.time()
+                time_elapsed=end_time-start_time
+                
+                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, Time elapsed: {:.4f}'
+                    .format(epoch, args.epochs, step, total_step, loss.item(), time_elapsed))
+                start_time=end_time
             step=step+1
 
         total_loss /= len(train_loader.dataset)
         train_score = 100 * train_score / len(train_loader.dataset)
 
+        print('Epoch [{}/{}], Training Loss: {:.4f}, Training Accuracy {:.4f}'
+                    .format(epoch, args.epochs, total_loss, train_score))
+        
+        model.train(False)
+        eval_score, bound, V_loss = evaluate_model(model, val_loader, device)
+        model.train(True)
+
+        logger.write('epoch %d, time: %.2f' % (epoch, time.time()-t))
+        logger.write('\ttrain_loss: %.3f, score: %.3f' % (total_loss, train_score))
+        logger.write('\teval loss: %.3f, score: %.3f (%.3f)' % (V_loss, 100 * eval_score, 100 * bound))
+
+        if eval_score > best_eval_score:
+            model_path = os.path.join(output_folder, 'model.pth')
+            torch.save(model.state_dict(), model_path)
+            best_eval_score = eval_score
+        
+        
+
 
 
 

Visual_Attention/utils.py

@@ -0,0 +1,100 @@
+from __future__ import print_function
+
+import errno
+import os
+import numpy as np
+from PIL import Image
+import torch
+import torch.nn as nn
+
+
+EPS = 1e-7
+
+
+def assert_eq(real, expected):
+    assert real == expected, '%s (true) vs %s (expected)' % (real, expected)
+
+
+def assert_array_eq(real, expected):
+    assert (np.abs(real-expected) < EPS).all(), \
+        '%s (true) vs %s (expected)' % (real, expected)
+
+
+def load_folder(folder, suffix):
+    imgs = []
+    for f in sorted(os.listdir(folder)):
+        if f.endswith(suffix):
+            imgs.append(os.path.join(folder, f))
+    return imgs
+
+
+def load_imageid(folder):
+    images = load_folder(folder, 'jpg')
+    img_ids = set()
+    for img in images:
+        img_id = int(img.split('/')[-1].split('.')[0].split('_')[-1])
+        img_ids.add(img_id)
+    return img_ids
+
+
+def pil_loader(path):
+    with open(path, 'rb') as f:
+        with Image.open(f) as img:
+            return img.convert('RGB')
+
+
+def weights_init(m):
+    """custom weights initialization."""
+    cname = m.__class__
+    if cname == nn.Linear or cname == nn.Conv2d or cname == nn.ConvTranspose2d:
+        m.weight.data.normal_(0.0, 0.02)
+    elif cname == nn.BatchNorm2d:
+        m.weight.data.normal_(1.0, 0.02)
+        m.bias.data.fill_(0)
+    else:
+        print('%s is not initialized.' % cname)
+
+
+def init_net(net, net_file):
+    if net_file:
+        net.load_state_dict(torch.load(net_file))
+    else:
+        net.apply(weights_init)
+
+
+def create_dir(path):
+    if not os.path.exists(path):
+        try:
+            os.makedirs(path)
+        except OSError as exc:
+            if exc.errno != errno.EEXIST:
+                raise
+
+
+class Logger(object):
+    def __init__(self, output_name):
+        dirname = os.path.dirname(output_name)
+        if not os.path.exists(dirname):
+            os.mkdir(dirname)
+
+        self.log_file = open(output_name, 'w')
+        self.infos = {}
+
+    def append(self, key, val):
+        vals = self.infos.setdefault(key, [])
+        vals.append(val)
+
+    def log(self, extra_msg=''):
+        msgs = [extra_msg]
+        for key, vals in self.infos.iteritems():
+            msgs.append('%s %.6f' % (key, np.mean(vals)))
+        msg = '\n'.join(msgs)
+        self.log_file.write(msg + '\n')
+        self.log_file.flush()
+        self.infos = {}
+        return msg
+
+    def write(self, msg):
+        self.log_file.write(msg + '\n')
+        self.log_file.flush()
+        print(msg)