def __init__(self, args, vocab_len):
super(BFM, self).__init__()
self.encoder = EncoderCNN(args.embed_size).eval().cpu()
self.encoder.load_state_dict(torch.load('encoder.ckpt', map_location=torch.device('cpu')))
self.decoder = DecoderRNN(args.embed_size, args.hidden_size, vocab_len, args.num_layers).eval().cpu()
self.decoder.forward = self.decoder.sample
self.decoder.load_state_dict(torch.load('decoder.ckpt', map_location=torch.device('cpu')))
def __init__(self):
print("Defining I.A")
# Device configuration
self.device = torch.device('cpu')
#vars
embed_size = 256
hidden_size = 512
num_layers = 1
encoder_path = 'models/encoder-5-3000.pkl'
decoder_path = 'models/decoder-5-3000.pkl'
vocab_path = 'data/vocab.pkl'
# Image preprocessing
self.transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])
with open(vocab_path, 'rb') as f:
self.vocab = pickle.load(f)
print("Building Model")
# Build models
self.encoder = EncoderCNN(embed_size).eval() # eval mode (batchnorm uses moving mean/variance)
self.decoder = DecoderRNN(embed_size, hidden_size, len(self.vocab), num_layers)
self.encoder = self.encoder.to(self.device)
self.decoder = self.decoder.to(self.device)
print("loading checkpoint")
# Load the trained model parameters
self.encoder.load_state_dict(torch.load(encoder_path))
self.decoder.load_state_dict(torch.load(decoder_path))
def main(args):
# Image preprocessing, normalization for the pretrained resnet
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225))])
# Load vocabulary wrapper
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
# Build models
encoder = EncoderCNN(args.embed_size).eval() # eval mode (batchnorm uses moving mean/variance)
decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), args.num_layers)
encoder = encoder.to(device)
decoder = decoder.to(device)
# Load the trained model parameters
encoder.load_state_dict(torch.load(args.encoder_path))
decoder.load_state_dict(torch.load(args.decoder_path))
# Build data loader
data_loader = get_loader(args.image_dir, args.caption_path, vocab,
transform, args.batch_size,
shuffle=True, num_workers=args.num_workers)
total_step = len(data_loader)
# List to score the BLEU scores
bleu_scores = []
for i, (images, captions, lengths) in enumerate(data_loader):
# Set mini-batch dataset
images = images.to(device)
# captions = captions.to(device)
# Generate an caption from the image
feature = encoder(images)
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids[0].cpu().numpy()
# Convert word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
score = sentence_bleu(captions, sentence, args.bleu_weights)
bleu_scores.append(score)
# Print log info
if i % args.log_step == 0:
print('Finish [{}/{}], Current BLEU Score: {:.4f}'
.format(i, total_step, np.mean(bleu_scores)))
np.save('test_results.npy', [bleu_scores, np.mean(bleu_scores)])
def test(args):
transform = transforms.Compose([
transforms.ToTensor(),
])
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
encoder = EncoderCNN(args.embed_size).eval()
decoder = DecoderRNN(args.embed_size, len(vocab), args.hidden_size,
args.num_layers)
# 加载训练好的模型的参数
encoder.load_state_dict(torch.load(args.encoder_path, map_location='cpu'))
decoder.load_state_dict(torch.load(args.decoder_path, map_location='cpu'))
image = load_img(args.img_path, transform)
feature = encoder(image)
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids[0].cpu().numpy()
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
print(sentence)
image = Image.open(args.img_path)
plt.imshow(np.asarray(image))
plt.show()
def getCaption(self,
imgs,
output_path='',
vocab_path='data/vocab.pkl',
decoder_path='models/decoder-5-3000.pkl',
encoder_path='models/encoder-5-3000.pkl',
embed_size=256,
hidden_size=512,
num_layers=1):
if (output_path == ''):
output_path = self.DEFAULT_OUTPUT_PATH
device = self.device
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
# Load vocabulary wrapper
with open(vocab_path, 'rb') as f:
vocab = pickle.load(f)
# Build models
encoder = EncoderCNN(embed_size).eval(
) # eval mode (batchnorm uses moving mean/variance)
decoder = DecoderRNN(embed_size, hidden_size, len(vocab), num_layers)
encoder = encoder.to(device)
decoder = decoder.to(device)
# Load the trained model parameters
encoder.load_state_dict(torch.load(encoder_path))
decoder.load_state_dict(torch.load(decoder_path))
CAPTIONS = []
for img in imgs:
# Prepare an image
image = self.load_image(img, transform=transform)
image_tensor = image.to(device)
# Generate an caption from the image
feature = encoder(image_tensor)
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids[0].cpu().numpy(
) # (1, max_seq_length) -> (max_seq_length)
# Convert word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
# Print out the image and the generated caption
CAPTIONS.append(self.prune_caption(sentence))
json_captions = self.writeJSON(imgs, CAPTIONS, output_path=output_path)
return json_captions
def __init__(self):
self.transform = transforms.Compose([
transforms.Resize(256), # smaller edge of image resized to 256
transforms.CenterCrop(224), # get 224x224 crop from the center
transforms.ToTensor(), # convert the PIL Image to a tensor
transforms.Normalize((0.485, 0.456, 0.406), # normalize image for pre-trained model
(0.229, 0.224, 0.225))])
# Load cherckpoint with best model
self.checkpoint = torch.load(os.path.join('./models', 'best-model.pkl'), 'cpu')
# Specify values for embed_size and hidden_size - we use the same values as in training step
self.embed_size = 512
self.hidden_size = 512
# Get the vocabulary and its size
self.vocab = Vocabulary(None, './vocab.pkl', "<start>", "<end>", "<unk>", "<pad>", "", "", True)
self.vocab_size = len(self.vocab)
# Initialize the encoder and decoder, and set each to inference mode
self.encoder = EncoderCNN(self.embed_size)
self.encoder.eval()
self.decoder = DecoderRNN(self.embed_size, self.hidden_size, self.vocab_size)
self.decoder.eval()
# Load the pre-trained weights
self.encoder.load_state_dict(self.checkpoint['encoder'])
self.decoder.load_state_dict(self.checkpoint['decoder'])
def main():
# Load vocabulary wrapper.
with open(vocab_path) as f:
vocab = pickle.load(f)
encoder = EncoderCNN(4096, embed_dim)
encoder.load_state_dict(torch.load('searchimage.pkl'))
for p in encoder.parameters():
p.requires_grad = False
word_encoder = EncoderRNN(embed_dim, embed_dim, len(vocab), num_layers_rnn)
word_encoder.load_state_dict(torch.load('searchword.pkl'))
if torch.cuda.is_available():
encoder.cuda()
word_encoder.cuda()
# Loss and Optimizer
criterion = nn.MSELoss()
params = list(
word_encoder.parameters()) # + list(encoder.linear.parameters())
optimizer = torch.optim.Adam(params, lr=2e-6, weight_decay=0.001)
#load data
with open(image_data_file) as f:
image_data = pickle.load(f)
image_features = si.loadmat(image_feature_file)
img_features = image_features['fc7'][0]
img_features = np.concatenate(img_features)
print 'here'
iteration = 0
for i in range(10): # epoch
use_caption = i % 5
print 'Epoch', i
losses = []
for x, y in make_mini_batch(img_features,
image_data,
use_caption=use_caption):
encoder.zero_grad()
word_encoder.zero_grad()
word_padding, lengths = make_word_padding(y, vocab)
x = Variable(torch.from_numpy(x).cuda())
word_index = Variable(torch.from_numpy(word_padding).cuda())
features = encoder(x)
outputs = word_encoder(word_index, lengths)
loss = torch.mean((features - outputs).pow(2))
loss.backward()
optimizer.step()
losses.append(loss.data[0])
if iteration % 100 == 0:
print 'loss', sum(losses) / float(len(losses))
losses = []
iteration += 1
torch.save(word_encoder.state_dict(), 'searchword.pkl')
torch.save(encoder.state_dict(), 'searchimage.pkl')
def run_inference(image_path,
encoder_path,
decoder_path,
vocab_path,
embed_size=256,
hidden_size=512,
num_layers=1):
print(f'sample.py running ... ')
# Image preprocessing
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
# Load vocabulary wrapper
with open(vocab_path, 'rb') as f:
print("using " + vocab_path)
vocab = pickle.load(f)
# Build models
encoder = EncoderCNN(
embed_size).eval() # eval mode (batchnorm uses moving mean/variance)
decoder = DecoderRNN(embed_size, hidden_size, len(vocab), num_layers)
encoder = encoder.to(device)
decoder = decoder.to(device)
# Load the trained model parameters
encoder.load_state_dict(
torch.load(encoder_path, map_location=torch.device('cpu')))
decoder.load_state_dict(
torch.load(decoder_path, map_location=torch.device('cpu')))
# Prepare an image
image = load_image(image_path, transform)
image_tensor = image.to(device)
# Generate an caption from the image
feature = encoder(image_tensor)
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids[0].cpu().numpy(
) # (1, max_seq_length) -> (max_seq_length)
# Convert word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
print(word)
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption).replace('<start>', '')
sentence = sentence.replace('<end>', '')
sentence = sentence.replace('_', ' ')
# Print out the image and the generated caption
print(sentence)
print(f'debug: chay xong roi ne')
return sentence.strip().capitalize()
def inference_coco(encoder_file: str, decoder_file: str, embed_size: int,
hidden_size: int, from_cpu: bool) -> None:
"""
Displays an original image from coco test dataset and prints its associated caption.
encoder_file: Name of the encoder to load.
decoder_file: Name of the decoder to load.
embed_size: Word embedding size for the encoder.
hidden_size: Hidden layer of the LSTM size.
from_cpu: Whether the model has been saved on CPU.
"""
# Define transform
transform_test = transforms.Compose([
transforms.Resize(256), # smaller edge of image resized to 256
transforms.RandomCrop(224), # get 224x224 crop from random location
transforms.ToTensor(), # convert the PIL Image to a tensor
transforms.Normalize(
(0.485, 0.456, 0.406), # normalize image for pre-trained model
(0.229, 0.224, 0.225))
])
# Device to use fo inference
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Create the data loader.
data_loader = get_loader(transform=transform_test, mode='test')
# Obtain sample image
_, image = next(iter(data_loader))
# The size of the vocabulary.
vocab_size = len(data_loader.dataset.vocab)
# Initialize the encoder and decoder, and set each to inference mode.
encoder = EncoderCNN(embed_size)
encoder.eval()
decoder = DecoderRNN(embed_size, hidden_size, vocab_size)
decoder.eval()
# Load the trained weights.
if from_cpu:
encoder.load_state_dict(
torch.load(os.path.join('./models', encoder_file),
map_location='cpu'))
decoder.load_state_dict(
torch.load(os.path.join('./models', decoder_file),
map_location='cpu'))
else:
encoder.load_state_dict(
torch.load(os.path.join('./models', encoder_file)))
decoder.load_state_dict(
torch.load(os.path.join('./models', decoder_file)))
# Move models to GPU if CUDA is available.
encoder.to(device)
decoder.to(device)
get_prediction(encoder, decoder, data_loader, device)
def main(args):
# Create model directory
if not os.path.exists(args.model_path):
os.makedirs(args.model_path)
# Build the models
encoder = EncoderCNN(args.embed_size).to(device)
decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab),
args.num_layers).to(device)
# Loss and optimizer
criterion = nn.CrossEntropyLoss()
params = list(decoder.parameters()) + list(
encoder.linear.parameters()) + list(encoder.bn.parameters())
optimizer = torch.optim.Adam(params, lr=args.learning_rate)
# Train the models
total_step = len(data_loader)
# For each TSP problem
for epoch in range(args.num_epochs):
for i, (images, captions, lengths) in enumerate(data_loader):
# Set mini-batch dataset
images = images.to(device)
captions = captions.to(device)
targets = pack_padded_sequence(captions, lengths,
batch_first=True)[0]
# Forward, backward and optimize
features = encoder(images)
outputs = decoder(features, captions, lengths)
loss = criterion(outputs, targets)
decoder.zero_grad()
encoder.zero_grad()
loss.backward()
optimizer.step()
# Print log info
if i % args.log_step == 0:
print(
'Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, Perplexity: {:5.4f}'
.format(epoch, args.num_epochs, i, total_step, loss.item(),
np.exp(loss.item())))
# Save the model checkpoints
if (i + 1) % args.save_step == 0:
torch.save(
decoder.state_dict(),
os.path.join(args.model_path,
'decoder-{}-{}.ckpt'.format(epoch + 1,
i + 1)))
torch.save(
encoder.state_dict(),
os.path.join(args.model_path,
'encoder-{}-{}.ckpt'.format(epoch + 1,
i + 1)))
def main(args):
# Image preprocessing
prediction = []
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225))])
# Load vocabulary wrapper
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
# Build models
encoder = EncoderCNN(args.embed_size).eval() # eval mode (batchnorm uses moving mean/variance)
decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), args.num_layers)
encoder = encoder.to(device)
decoder = decoder.to(device)
# Load the trained model parameters
encoder.load_state_dict(torch.load(args.encoder_path))
decoder.load_state_dict(torch.load(args.decoder_path))
dirname=''
fnames = listdir(os.getcwd)
#with open(dirname)
for fname in fnames:
#print(fname)
# Prepare an image
image = load_image(''+fname, transform)
image_tensor = image.to(device)
# Generate an caption from the image
feature = encoder(image_tensor)
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids[0].cpu().numpy() # (1, max_seq_length) -> (max_seq_length)
# Convert word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
current_pred = [fname, sentence]
predictions.append(current_pred)
# Print out the image and the generated caption
print(fname)
print (sentence)
#image = Image.open(args.image)
#plt.imshow(np.asarray(image))
df = pd.DataFrame(predictions, columns=['File Name', 'Caption'])
df.to_excel('output.xls')
class Neuraltalk2:
def __init__(self):
print("Defining I.A")
# Device configuration
self.device = torch.device('cpu')
#vars
embed_size = 256
hidden_size = 512
num_layers = 1
encoder_path = 'models/encoder-5-3000.pkl'
decoder_path = 'models/decoder-5-3000.pkl'
vocab_path = 'data/vocab.pkl'
# Image preprocessing
self.transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])
with open(vocab_path, 'rb') as f:
self.vocab = pickle.load(f)
print("Building Model")
# Build models
self.encoder = EncoderCNN(embed_size).eval() # eval mode (batchnorm uses moving mean/variance)
self.decoder = DecoderRNN(embed_size, hidden_size, len(self.vocab), num_layers)
self.encoder = self.encoder.to(self.device)
self.decoder = self.decoder.to(self.device)
print("loading checkpoint")
# Load the trained model parameters
self.encoder.load_state_dict(torch.load(encoder_path))
self.decoder.load_state_dict(torch.load(decoder_path))
def eval_image(self, image_path):
# Prepare an image
image = load_image(image_path, self.transform)
image_tensor = image.to(self.device)
# Generate an caption from the image
feature = self.encoder(image_tensor)
sampled_ids = self.decoder.sample(feature)
sampled_ids = sampled_ids[0].cpu().numpy()
# Convert word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = self.vocab.idx2word[word_id]
if word == '<end>':
break
if word == '<start>':
continue
sampled_caption.append(word)
sentence = ' '.join(sampled_caption)
return sentence
def predict(self, args):
print('predict..start')
device = torch.device('cpu')
# Image preprocessing
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
# Load vocabulary wrapper
# Load vocabulary wrapper
#vocab = Vocabulary()
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
# Build models
encoder = EncoderCNN(args.embed_size).eval(
) # eval mode (batchnorm uses moving mean/variance)
decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab),
args.num_layers)
encoder = encoder.to(device)
decoder = decoder.to(device)
# Load the trained model parameters
encoder.load_state_dict(
torch.load(args.encoder_path, map_location=device))
decoder.load_state_dict(
torch.load(args.decoder_path, map_location=device))
# Prepare an image
image = self.load_image(args.image, transform)
image_tensor = image.to(device)
# Generate an caption from the image
feature = encoder(image_tensor)
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids[0].cpu().numpy(
) # (1, max_seq_length) -> (max_seq_length)
# Convert word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
# Print out the image and the generated caption
print(sentence)
return sentence
def main(cfg):
# print(cfg.pretty())
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(cfg.image.mean, cfg.image.std)
])
print(hydra.utils.to_absolute_path(cfg.train.vocab_path))
with open(hydra.utils.to_absolute_path(cfg.train.vocab_path), 'rb') as f:
vocab = pickle.load(f)
# モデルの構築
encoder = EncoderCNN(cfg.train.embed_size).eval()
decoder = DecoderRNN(cfg.train.embed_size, cfg.train.hidden_size,
len(vocab), cfg.train.num_layers)
encoder = encoder.to(device)
decoder = decoder.to(device)
# 学習済みモデルのパラメータを読み込む
encoder.load_state_dict(
torch.load(hydra.utils.to_absolute_path(cfg.train.encoder_path)))
decoder.load_state_dict(
torch.load(hydra.utils.to_absolute_path(cfg.train.decoder_path)))
with open(json_dir, encoding='utf-8') as f:
data = json.loads(f.read())
for key in data['images']:
img_file_name = key['file_name']
img_file_path = base_dir + '/data/val2014/' + img_file_name
# 画像の準備
image = load_image(hydra.utils.to_absolute_path(img_file_path),
transform)
image_tensor = image.to(device)
# 入力した画像からキャプションを生成する
feature = encoder(image_tensor)
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids[0].cpu().numpy()
# word_idsをwordに変換する
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
# <start>,<end>,"."を取り除いて処理を行う
sentence = ' '.join(sampled_caption[1:-2])
print(sentence)
def main2(image,
encoder_path='models/encoder-5-3000.pkl',
decoder_path='models/decoder-5-3000.pkl',
vocab_path="data/vocab.pkl",
embed_size=256,
hidden_size=512,
num_layers=1):
# Image preprocessing
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
# Load vocabulary wrapper
with open(vocab_path, 'rb') as f:
vocab = pickle.load(f)
# Build models
encoder = EncoderCNN(
embed_size).eval() # eval mode (batchnorm uses moving mean/variance)
decoder = DecoderRNN(embed_size, hidden_size, len(vocab), num_layers)
encoder = encoder.to(device)
decoder = decoder.to(device)
# Load the trained model parameters
encoder.load_state_dict(torch.load(encoder_path))
decoder.load_state_dict(torch.load(decoder_path))
# Prepare an image
image = load_image(image, transform)
image_tensor = image.to(device)
# Generate an caption from the image
feature = encoder(image_tensor)
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids[0].cpu().numpy(
) # (1, max_seq_length) -> (max_seq_length)
# Convert word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
# Print out the image and the generated caption
print(sentence)
#image = Image.open(args.image)
#plt.imshow(np.asarray(image))
return sentence
class BFM(nn.Module):
def __init__(self, args, vocab_len):
super(BFM, self).__init__()
self.encoder = EncoderCNN(args.embed_size).eval().cpu()
self.encoder.load_state_dict(torch.load('encoder.ckpt', map_location=torch.device('cpu')))
self.decoder = DecoderRNN(args.embed_size, args.hidden_size, vocab_len, args.num_layers).eval().cpu()
self.decoder.forward = self.decoder.sample
self.decoder.load_state_dict(torch.load('decoder.ckpt', map_location=torch.device('cpu')))
def forward(self, image):
feature = self.encoder(image)
sampled_ids = self. decoder(feature)
return sampled_ids
def main(args):
# Image preprocessing
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
# Load vocabulary wrapper
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
# Build models
encoder = EncoderCNN(args.embed_size).eval(
) # eval mode (batchnorm uses moving mean/variance)
decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab),
args.num_layers)
encoder = encoder.to(device)
decoder = decoder.to(device)
# Load the trained model parameters
encoder.load_state_dict(torch.load(args.encoder_path))
decoder.load_state_dict(torch.load(args.decoder_path))
# Prepare an image
image = load_image(args.image, transform)
image_tensor = image.to(device)
# Generate an caption from the image
feature = encoder(image_tensor)
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids[0].cpu().numpy(
) # (1, max_seq_length) -> (max_seq_length)
# Convert word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sampled_caption.remove('<end>')
sampled_caption.remove('<start>')
sentence = ' '.join(sampled_caption)
# Print out the image and the generated caption
print(sentence)
f = open("demofile3.txt", "w")
f.truncate(0)
f.write(sentence)
f.close()
image = Image.open(args.image)
plt.imshow(np.asarray(image))
def main():
# TODO parse arguments
parser = argparse.ArgumentParser()
# DATA related
parser.add_argument('--data_path', type=str,
help='Path of training data')
parser.add_argument('--max_length', type=int, default=250,
help='Pad all data to this size')
# Model Specs
parser.add_argument('--embed_size', type=int, default=300,
help='Dimension of word embedding')
parser.add_argument('--filter_size', type=int, default=5,
help='size of filter')
parser.add_argument('--stride', type=int, default=2,
help='stride size for each layer')
parser.add_argument('--filter_nums', type=str, default='300,600',
help='filer number for each convolution layer')
parser.add_argument('--hidden_size', type=int, default=500,
help='size of hidden state in the middle')
# Optimization Specs
parser.add_argument('--batch_size', type=int, default=128,
help='Minibatch during training')
parser.add_argument('--optim', type=str, default='SGD',
help='Optimization method')
parser.add_argument('--learning_rate', type=float, default=0.0005,
help='Learning rate of model')
parser.add_argument('--decay', type=float, default=0.5,
help='Decay rate of learning rate')
parser.add_argument('--start_decay', type=int, default=5,
help='Start Epoch of decay learning rate')
args = parser.parse_args()
args.filter_nums = [int(i) for i in args.filter_nums.split(',')]
args.filter_nums.append(args.hidden_size)
final_filter_size = args.max_length
for num in args.filter_nums[:-1]:
final_filter_size = math.floor(
(final_filter_size - args.filter_size)/args.stride + 1
)
# TODO build model
encoder = EncoderCNN(args.embed_size, args.filter_size,
args.stride, args.filter_nums, final_filter_size)
decoder = DecoderCNN(args.embed_size, args.filter_size,
args.stride, args.filter_nums, final_filter_size)
if torch.cuda.is_available():
encoder.cuda()
decoder.cuda()
creterion =
# TODO Load training data
with open(args.data_path, 'r') as f:
data = f.readlines()
def main(args):
with open('data/vocab.pkl', 'rb') as f:
vocab = pickle.load(f)
encoder = EncoderCNN(256)
encoder.eval() # evaluation mode (BN uses moving mean/variance)
decoder = DecoderRNN(256, 512, len(vocab), 1)
if torch.cuda.is_available():
encoder.cuda()
decoder.cuda()
# Load the trained model parameters
encoder.load_state_dict(torch.load(args.encoder))
decoder.load_state_dict(torch.load(args.decoder))
measurement_score = test(encoder, decoder, vocab, args.num_samples,
args.num_hints, args.debug, args.c_step,
args.no_avg)
if args.msm == "co":
scores = cocoEval()
scores_u = cocoEval(res='data/captions_val2014_results_u.json')
print(scores)
print(scores_u)
with open(args.filepath, 'w+') as f:
pickle.dump((scores, scores_u), f)
def main(args):
# Build models
encoder = EncoderCNN(args.embed_size).eval(
) # eval mode (batchnorm uses moving mean/variance)
decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab),
args.num_layers)
encoder = encoder.to(device)
decoder = decoder.to(device)
# Load the trained model parameters
encoder.load_state_dict(torch.load(
args.encoder_path)) #path for dumbing output of encoder model
decoder.load_state_dict(torch.load(
args.decoder_path)) #path for dumbing output of decoder model
test_data = generate_training_data(5)
textures_test = generate_textures(test_data)
transforms_test = generate_transforms(test_data)
for i in range(len(textures_test)):
plt.imsave('predictions/texture4_0%i.png' % i,
textures_test[i],
cmap="gray")
print(transforms_test)
predicted_progs = []
for texture in textures_test:
texture = torch.tensor(texture, device=device)
texture = texture.unsqueeze(0)
texture = texture.unsqueeze(
0) #for EncoderCNN ought to unsqueeze twice
feature = encoder(texture)
sampled_seq = decoder.sample(feature)
sampled_seq = sampled_seq[0].cpu().numpy(
) # (1, max_seq_length) -> (max_seq_length)
# Convert sampled sequence of transforms to words
prog = []
for int_word in sampled_seq:
word = int_to_word(int_word)
prog.append(word)
if word == '<end>':
break
trans_seq = '-->'.join(prog)
predicted_progs.append([trans_seq])
# Print out the sequence of generated transform sequences
print(predicted_progs)
def main(args):
# Val images folder
filepath = '/scratch/ys2542/pytorch-tutorial/tutorials/03-advanced/image_captioning/data/resizedval2014'
onlyfiles = [fl for fl in listdir(filepath) if isfile(join(filepath, fl))]
# image preprocessing
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
# load vocabulary wrapper pickle file
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
encoder = EncoderCNN(args.embed_size) # build encoder
encoder.eval() # evaluation mode by moving mean and variance
decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab),
args.num_layers) # build decoder
# load the trained CNN and RNN parameters
encoder.load_state_dict(torch.load(args.encoder_path))
decoder.load_state_dict(torch.load(args.decoder_path))
# Load all images in val folder
for i in onlyfiles:
badsize = 0 # count the unload images
args_image = filepath + '/' # val folder path with image names
args_image = args_image + i
# transform image and wrap it to tensor
image = load_image(args_image, transform)
image_tensor = to_var(image, volatile=True)
if torch.cuda.is_available(): # load GPU
encoder.cuda()
decoder.cuda()
# generate caption from image
try:
feature = encoder(image_tensor)
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids.cpu().data.numpy()
# decode word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
# print out image and generated caption without start and end
print('beam_size_1' + '\t' + i + '\t' + sentence[8:-8])
except:
badsize = badsize + 1 # count some wrong images
def main(image):
# Configuration for hyper-parameters
config = Config()
# Image Preprocessing
transform = config.test_transform
# Load vocabulary
with open(os.path.join(config.vocab_path, 'vocab.pkl'), 'rb') as f:
vocab = pickle.load(f)
# Build Models
encoder = EncoderCNN(config.embed_size)
encoder.eval() # evaluation mode (BN uses moving mean/variance)
decoder = DecoderRNN(config.embed_size, config.hidden_size, len(vocab),
config.num_layers)
# Load the trained model parameters
encoder.load_state_dict(
torch.load(
os.path.join(config.teacher_cnn_path, config.trained_encoder)))
decoder.load_state_dict(
torch.load(
os.path.join(config.teacher_lstm_path, config.trained_decoder)))
# Prepare Image
image = Image.open(image)
image_tensor = Variable(transform(image).unsqueeze(0))
# Set initial states
state = (Variable(torch.zeros(config.num_layers, 1, config.hidden_size)),
Variable(torch.zeros(config.num_layers, 1, config.hidden_size)))
# If use gpu
if torch.cuda.is_available():
encoder.cuda()
decoder.cuda()
state = [s.cuda() for s in state]
image_tensor = image_tensor.cuda()
# Generate caption from image
feature = encoder(image_tensor)
sampled_ids = decoder.sample(feature, state)
sampled_ids = sampled_ids.cpu().data.numpy()
# Decode word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word_id == 96:
sampled_caption.append('<end>')
break
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
# Print out image and generated caption.
print(sentence)
return sentence
def main(args):
# Image preprocessing
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
# Load vocabulary wrapper
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
# Build Models
encoder = EncoderCNN(args.embed_size)
encoder.eval() # evaluation mode (BN uses moving mean/variance)
decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab),
args.num_layers)
# Load the trained model parameters
encoder.load_state_dict(torch.load(args.encoder_path))
decoder.load_state_dict(torch.load(args.decoder_path))
# Prepare Image
image = load_image(args.image, transform)
image_tensor = to_var(image, volatile=True)
# If use gpu
if torch.cuda.is_available():
encoder.cuda()
decoder.cuda()
# Generate caption from image
feature = encoder(image_tensor)
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids.cpu().data.numpy()
# Decode word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
# Print out image and generated caption.
sentence = sentence.replace('<start> ',
'').replace(' <end>', '').replace('.',
'').strip()
translator = Translator()
sentence_indo = translator.translate(sentence, dest='id').text
print('This is an image of: ' + sentence_indo)
tts = gTTS(sentence_indo, 'id')
tts.save('result.mp3')
playsound('result.mp3')
image = Image.open(args.image)
plt.imshow(np.asarray(image))
plt.show()
def main(args):
# Image preprocessing
transform = transforms.Compose([
transforms.Scale(args.crop_size),
transforms.CenterCrop(args.crop_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# Load vocabulary wrapper
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
alexnet = models.alexnet(pretrained=True)
alexnet2 = AlexNet2(alexnet)
# Build Models
encoder = EncoderCNN(4096, args.embed_size)
encoder.eval() # evaluation mode (BN uses moving mean/variance)
decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab),
args.num_layers)
# Load the trained model parameters
encoder.load_state_dict(torch.load(args.encoder_path))
decoder.load_state_dict(torch.load(args.decoder_path))
# Prepare Image
image = Image.open(args.image)
image_tensor = Variable(transform(image).unsqueeze(0))
# Set initial states
state = (Variable(torch.zeros(args.num_layers, 1, args.hidden_size)),
Variable(torch.zeros(args.num_layers, 1, args.hidden_size)))
# If use gpu
if torch.cuda.is_available():
encoder.cuda()
decoder.cuda()
alexnet2.cuda()
state = [s.cuda() for s in state]
image_tensor = image_tensor.cuda()
# Generate caption from image
alexnet2(image_tensor)
feature = encoder(alexnet2.fc7_value)
sampled_ids = decoder.sample(feature, state)
sampled_ids = sampled_ids.cpu().data.numpy()
# Decode word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
# Print out image and generated caption.
print(sentence)
def main(args):
# Image preprocessing
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225))])
# Load vocabulary wrapper
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
# Build models
encoder = EncoderCNN(args.embed_size).eval() # eval mode (batchnorm uses moving mean/variance)
decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), args.num_layers)
encoder = encoder.to(device)
decoder = decoder.to(device)
# Load the trained model parameters
encoder.load_state_dict(torch.load(args.encoder_path))
decoder.load_state_dict(torch.load(args.decoder_path))
# Prepare an image
'''
image = load_image(args.image, transform)
image_tensor = image.to(device)
'''
data_loader, _ = get_loader(transforms=False)
inp, targets = next(iter(data_loader))
audio = inp_transform(inp)
audio = audio.to(device)
# Generate an caption from the image
feature = encoder(audio)
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids[0].cpu().numpy() # (1, max_seq_length) -> (max_seq_length)
# Convert word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
# Print out the image and the generated caption
print("Logits : {}\nTarget : {}".format(sentence, targets))
'''
def get_encoder_decoder(vocab):
""" Given the arguments, returns the correct combination of CNN/RNN/GAN encoders and decoders. """
if args.pretrain_rnn:
encoder = EncoderRNN(len(vocab),
args.embed_size,
args.encoder_rnn_hidden_size,
num_layers=args.num_layers).to(device)
elif args.gan_embedding:
gan = torch.load('DCGAN_embed_2.tch').to(device)
encoder = gan.discriminator
elif args.progan_embedding:
pro_gan = pg.ProGAN(depth=7,
latent_size=256,
device=torch.device('cuda'))
pro_gan.dis.load_state_dict(torch.load('progan_weights/GAN_DIS_6.pth'))
# pro_gan.dis_optim.load_state_dict(torch.load('progan_weights/GAN_DIS_OPTIM_6.pth'))
pro_gan.gen.load_state_dict(torch.load('progan_weights/GAN_GEN_6.pth'))
# pro_gan.gen_optim.load_state_dict(torch.load('progan_weights/GAN_GEN_OPTIM_6.pth'))
pro_gan.gen_shadow.load_state_dict(
torch.load('progan_weights/GAN_GEN_SHADOW_6.pth'))
print("Loaded proGAN weights.", flush=True)
encoder = pro_gan.dis.to(device)
else:
encoder = EncoderCNN(args.embed_size).to(device)
decoder = DecoderRNNOld(args.embed_size,
args.decoder_rnn_hidden_size,
len(vocab),
args.num_layers,
vocab,
device=device).to(device)
return encoder, decoder
def main(args):
# Image preprocessing
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
# Load vocabulary wrapper
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
# Build Models
encoder = EncoderCNN(args.embed_size)
encoder.eval() # evaluation mode (BN uses moving mean/variance)
decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab),
args.num_layers)
# Load the trained model parameters
encoder.load_state_dict(torch.load(args.encoder_path))
decoder.load_state_dict(torch.load(args.decoder_path))
# Prepare Image
#image = load_image(args.image, transform)
#image_tensor = to_var(image, volatile=True)
# If use gpu
if torch.cuda.is_available():
encoder.cuda()
decoder.cuda()
data = []
try:
img_path = args.image
# Prepare Image
image = load_image(img_path, transform)
image_tensor = to_var(image, volatile=True)
# Generate caption from image
feature = encoder(image_tensor)
#pdb.set_trace()
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids.cpu().data.numpy()
# Decode word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
if word == '<start>':
continue
if word == '<end>':
break
sampled_caption.append(word)
sentence = ' '.join(sampled_caption)
# Print out image and generated caption.
print(sentence)
data.append({'key': img_path.split('/')[-1], 'sentence': sentence})
except:
print(img_path)
def main(args):
# Image preprocessing
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225))])
# Load vocabulary wrapper
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
# Build Models
encoder = EncoderCNN(args.embed_size)
encoder.eval() # evaluation mode (BN uses moving mean/variance)
decoder = DecoderRNN(args.embed_size, args.hidden_size,
len(vocab), args.num_layers)
# Load the trained model parameters
encoder.load_state_dict(torch.load(args.encoder_path))
decoder.load_state_dict(torch.load(args.decoder_path))
# If use gpu
if torch.cuda.is_available():
encoder.cuda()
decoder.cuda()
# Prepare Image
image_dir = args.image
images = os.listdir(image_dir)
for image_id in images:
if not image_id.endswith('.jpg'):
continue
image = os.path.join(image_dir, image_id)
image = load_image(image, transform)
image_tensor = to_var(image, volatile=True)
# Generate caption from image
try:
feature, cnn_features = encoder(image_tensor)
sampled_ids = decoder.sample(feature, cnn_features)
sampled_ids = sampled_ids.cpu().data.numpy()
except:
continue
# Decode word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
# Print out image and generated caption.
print (image_id + '\t' + sentence)
def main(args):
# Image preprocessing
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
# Load vocabulary wrapper
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
# Build models
encoder = EncoderCNN(args.embed_size).eval(
) # eval mode (batchnorm uses moving mean/variance)
decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab),
args.num_layers)
encoder = encoder.to(device)
decoder = decoder.to(device)
# Load the trained model parameters
encoder.load_state_dict(torch.load(args.encoder_path))
decoder.load_state_dict(torch.load(args.decoder_path))
images = glob.glob('../data/flickr/train/*.jpg')
for i in images:
# Prepare an image
image = load_image(i, transform)
image_tensor = image.to(device)
# Generate an caption from the image
feature = encoder(image_tensor)
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids[0].cpu().numpy(
) # (1, max_seq_length) -> (max_seq_length)
# Convert word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
print("{} : {}".format(i[21:], sentence))
def main(cfg):
# print(cfg.pretty())
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(cfg.image.mean, cfg.image.std)
])
print(hydra.utils.to_absolute_path(cfg.train.vocab_path))
with open(hydra.utils.to_absolute_path(cfg.train.vocab_path), 'rb') as f:
vocab = pickle.load(f)
# モデルの構築
encoder = EncoderCNN(cfg.train.embed_size).eval()
decoder = DecoderRNN(cfg.train.embed_size, cfg.train.hidden_size,
len(vocab), cfg.train.num_layers)
encoder = encoder.to(device)
decoder = decoder.to(device)
# 学習済みモデルのパラメータを読み込む
encoder.load_state_dict(
torch.load(hydra.utils.to_absolute_path(cfg.train.encoder_path)))
decoder.load_state_dict(
torch.load(hydra.utils.to_absolute_path(cfg.train.decoder_path)))
# 画像の準備
image = load_image(hydra.utils.to_absolute_path(cfg.sample.image_path),
transform)
image_tensor = image.to(device)
# 入力した画像からキャプションを生成する
feature = encoder(image_tensor)
# sampled_ids = decoder.sample(feature)
sampled_ids = decoder.greedy_decode(features=feature)
sampled_ids = sampled_ids[0].cpu().numpy()
# word_idsをwordに変換する
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
print(sentence)
image = Image.open(hydra.utils.to_absolute_path(cfg.sample.image_path))
plt.imshow(np.asarray(image))
def main(args):
# Image preprocessing
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225))])
# Load vocabulary wrapper
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
# Build models
encoder = EncoderCNN(args.embed_size).eval() # eval mode (batchnorm uses moving mean/variance)
decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), args.num_layers)
encoder = encoder.to(device)
decoder = decoder.to(device)
# Load the trained model parameters
encoder.load_state_dict(torch.load(args.encoder_path))
decoder.load_state_dict(torch.load(args.decoder_path))
# Prepare an image
image = load_image(args.image, transform)
image_tensor = image.to(device)
# Generate an caption from the image
feature = encoder(image_tensor)
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids[0].cpu().numpy() # (1, max_seq_length) -> (max_seq_length)
# Convert word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
# Print out the image and the generated caption
print (sentence)
image = Image.open(args.image)
plt.imshow(np.asarray(image))
def main(args):
# Image preprocessing
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225))])
# Load vocabulary wrapper
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
# Build Models
encoder = EncoderCNN(args.embed_size)
encoder.eval() # evaluation mode (BN uses moving mean/variance)
decoder = DecoderRNN(args.embed_size, args.hidden_size,
len(vocab), args.num_layers)
# Load the trained model parameters
encoder.load_state_dict(torch.load(args.encoder_path))
decoder.load_state_dict(torch.load(args.decoder_path))
# Prepare Image
image = load_image(args.image, transform)
image_tensor = to_var(image, volatile=True)
# If use gpu
if torch.cuda.is_available():
encoder.cuda()
decoder.cuda()
# Generate caption from image
feature = encoder(image_tensor)
sampled_ids = decoder.sample(feature)
sampled_ids = sampled_ids.cpu().data.numpy()
# Decode word_ids to words
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sampled_caption)
# Print out image and generated caption.
print (sentence)
image = Image.open(args.image)
plt.imshow(np.asarray(image))
def main(args):
# Create model directory
if not os.path.exists(args.model_path):
os.makedirs(args.model_path)
# Image preprocessing, normalization for the pretrained resnet
transform = transforms.Compose([
transforms.RandomCrop(args.crop_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225))])
# Load vocabulary wrapper
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
# Build data loader
data_loader = get_loader(args.image_dir, args.caption_path, vocab,
transform, args.batch_size,
shuffle=True, num_workers=args.num_workers)
# Build the models
encoder = EncoderCNN(args.embed_size).to(device)
decoder = DecoderRNN(args.embed_size, args.hidden_size, len(vocab), args.num_layers).to(device)
# Loss and optimizer
criterion = nn.CrossEntropyLoss()
params = list(decoder.parameters()) + list(encoder.linear.parameters()) + list(encoder.bn.parameters())
optimizer = torch.optim.Adam(params, lr=args.learning_rate)
# Train the models
total_step = len(data_loader)
for epoch in range(args.num_epochs):
for i, (images, captions, lengths) in enumerate(data_loader):
# Set mini-batch dataset
images = images.to(device)
captions = captions.to(device)
targets = pack_padded_sequence(captions, lengths, batch_first=True)[0]
# Forward, backward and optimize
features = encoder(images)
outputs = decoder(features, captions, lengths)
loss = criterion(outputs, targets)
decoder.zero_grad()
encoder.zero_grad()
loss.backward()
optimizer.step()
# Print log info
if i % args.log_step == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, Perplexity: {:5.4f}'
.format(epoch, args.num_epochs, i, total_step, loss.item(), np.exp(loss.item())))
# Save the model checkpoints
if (i+1) % args.save_step == 0:
torch.save(decoder.state_dict(), os.path.join(
args.model_path, 'decoder-{}-{}.ckpt'.format(epoch+1, i+1)))
torch.save(encoder.state_dict(), os.path.join(
args.model_path, 'encoder-{}-{}.ckpt'.format(epoch+1, i+1)))
transforms.Normalize((0.485, 0.456, 0.406), # normalize image for pre-trained model
(0.229, 0.224, 0.225))])
# Build data loader.
data_loader = get_loader(transform=transform_train,
mode='train',
batch_size=batch_size,
vocab_threshold=vocab_threshold,
vocab_from_file=vocab_from_file,
cocoapi_loc=COCOPATH)
# The size of the vocabulary.
vocab_size = len(data_loader.dataset.vocab)
# Initialize the encoder and decoder.
encoder = EncoderCNN(embed_size)
decoder = DecoderRNN(embed_size, hidden_size, vocab_size)
# Move models to GPU if CUDA is available.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
encoder.to(device)
decoder.to(device)
# Define the loss function.
criterion = nn.CrossEntropyLoss().cuda() if torch.cuda.is_available() else nn.CrossEntropyLoss()
# TODO #3: Specify the learnable parameters of the model.
params = list(decoder.parameters()) +\
list(encoder.embed.parameters()) # We don't want to retrain the resnet
# TODO #4: Define the optimizer.
(0.229, 0.224, 0.225))])
# Build data loader.
data_loader = get_loader(transform=transform_test,
mode='test_small',
batch_size=batch_size,
vocab_threshold=vocab_threshold,
vocab_from_file=vocab_from_file,
cocoapi_loc=COCOPATH)
vocab = data_loader.dataset.vocab
# The size of the vocabulary.
vocab_size = len(vocab)
# Initialize the encoder and decoder.
encoder = EncoderCNN(embed_size)
decoder = DecoderRNN(embed_size, hidden_size, vocab_size)
device = torch.device("cpu")
# encoder.to(device)
# decoder.to(device)
# Load the pretrained model
encoder.load_state_dict(torch.load(PRETRAINED_MODEL_PATH.format('encoder')))
decoder.load_state_dict(torch.load(PRETRAINED_MODEL_PATH.format('decoder')))
encoder.eval()
decoder.eval()
images, conv_images = next(iter(data_loader))
features = encoder(conv_images).unsqueeze(1)
def main(args):
# Create model directory
if not os.path.exists(args.model_path):
os.makedirs(args.model_path)
# Image preprocessing
# For normalization, see https://github.com/pytorch/vision#models
transform = transforms.Compose([
transforms.RandomCrop(args.crop_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225))])
# Load vocabulary wrapper.
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
# Build data loader
data_loader = get_loader(args.image_dir, args.caption_path, vocab,
transform, args.batch_size,
shuffle=True, num_workers=args.num_workers)
# Build the models
encoder = EncoderCNN(args.embed_size)
decoder = DecoderRNN(args.embed_size, args.hidden_size,
len(vocab), args.num_layers)
if torch.cuda.is_available():
encoder.cuda()
decoder.cuda()
# Loss and Optimizer
criterion = nn.CrossEntropyLoss()
params = list(decoder.parameters()) + list(encoder.linear.parameters()) + list(encoder.bn.parameters())
optimizer = torch.optim.Adam(params, lr=args.learning_rate)
# Train the Models
total_step = len(data_loader)
for epoch in range(args.num_epochs):
for i, (images, captions, lengths) in enumerate(data_loader):
# Set mini-batch dataset
images = to_var(images, volatile=True)
captions = to_var(captions)
targets = pack_padded_sequence(captions, lengths, batch_first=True)[0]
# Forward, Backward and Optimize
decoder.zero_grad()
encoder.zero_grad()
features = encoder(images)
outputs = decoder(features, captions, lengths)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
# Print log info
if i % args.log_step == 0:
print('Epoch [%d/%d], Step [%d/%d], Loss: %.4f, Perplexity: %5.4f'
%(epoch, args.num_epochs, i, total_step,
loss.data[0], np.exp(loss.data[0])))
# Save the models
if (i+1) % args.save_step == 0:
torch.save(decoder.state_dict(),
os.path.join(args.model_path,
'decoder-%d-%d.pkl' %(epoch+1, i+1)))
torch.save(encoder.state_dict(),
os.path.join(args.model_path,
'encoder-%d-%d.pkl' %(epoch+1, i+1)))
