Update dataloader raw.

Author: ruotianluo

README.md

@@ -6,7 +6,6 @@ There's something difference compared to neuraltalk2.
 - Use resnet101; the same way as in self-critical (the preprocessing code may have bug, haven't tested yet)
 
 # TODO:
-- eval code for arbitrary images
 - Other models
 
 # Requirements

dataloaderraw.py

@@ -7,10 +7,27 @@
 import os
 import numpy as np
 import random
+import torch
+from torch.autograd import Variable
 import skimage
 import skimage.io
 import scipy.misc
 
+from torchvision import transforms as trn
+preprocess = trn.Compose([
+        #trn.ToTensor(),
+        trn.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+from misc.resnet_utils import myResnet
+import misc.resnet as resnet
+
+resnet = resnet.resnet101()
+resnet.load_state_dict(torch.load('/home-nfs/rluo/rluo/model/pytorch-resnet/resnet101.pth'))
+my_resnet = myResnet(resnet)
+my_resnet.cuda()
+my_resnet.eval()
+
 class DataLoaderRaw():
 
     def __init__(self, opt):
@@ -65,7 +82,8 @@ def get_batch(self, split, batch_size=None):
         batch_size = batch_size or self.batch_size
 
         # pick an index of the datapoint to load next
-        img_batch = np.ndarray([batch_size, 224,224,3], dtype = 'float32')
+        fc_batch = np.ndarray((batch_size, 2048), dtype = 'float32')
+        att_batch = np.ndarray((batch_size, 14, 14, 2048), dtype = 'float32')
         max_index = self.N
         wrapped = False
         infos = []
@@ -85,15 +103,22 @@ def get_batch(self, split, batch_size=None):
                 img = img[:,:,np.newaxis]
                 img = img.concatenate((img, img, img), axis=2)
 
-            img_batch[i] = img[16:240, 16:240, :].astype('float32')/255.0
+            img = img.astype('float32')/255.0
+            img = torch.from_numpy(img.transpose([2,0,1])).cuda()
+            img = Variable(preprocess(img), volatile=True)
+            tmp_fc, tmp_att = my_resnet(img)
+
+            fc_batch[i] = tmp_fc.data.cpu().float().numpy()
+            att_batch[i] = tmp_att.data.cpu().float().numpy()
 
             info_struct = {}
             info_struct['id'] = self.ids[ri]
             info_struct['file_path'] = self.files[ri]
             infos.append(info_struct)
 
         data = {}
-        data['images'] = img_batch
+        data['fc_feats'] = fc_batch
+        data['att_feats'] = att_batch
         data['bounds'] = {'it_pos_now': self.iterator, 'it_max': self.N, 'wrapped': wrapped}
         data['infos'] = infos
 

train.py

@@ -21,8 +21,6 @@
 
 import os
 
-#from ipdb import set_trace
-
 def train(opt):
     loader = DataLoader(opt)
     opt.vocab_size = loader.vocab_size

train_tb.py

@@ -22,11 +22,6 @@
 import misc.utils as utils
 import tensorflow as tf
 
-import os
-NUM_THREADS = 2 #int(os.environ['OMP_NUM_THREADS'])
-
-#from ipdb import set_trace
-
 def add_summary_value(writer, key, value, iteration):
     summary = tf.Summary(value=[tf.Summary.Value(tag=key, simple_value=value)])
     writer.add_summary(summary, iteration)