Skip to content

Commit

Permalink
Update dataloader raw.
Browse files Browse the repository at this point in the history
  • Loading branch information
ruotianluo committed May 3, 2017
1 parent 503d2f5 commit 397f1a8
Show file tree
Hide file tree
Showing 4 changed files with 28 additions and 11 deletions.
1 change: 0 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
31 changes: 28 additions & 3 deletions dataloaderraw.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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 = []
Expand All @@ -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

Expand Down
2 changes: 0 additions & 2 deletions train.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,6 @@

import os

#from ipdb import set_trace

def train(opt):
loader = DataLoader(opt)
opt.vocab_size = loader.vocab_size
Expand Down
5 changes: 0 additions & 5 deletions train_tb.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit 397f1a8

Please sign in to comment.