algorithm/pytorch-tutorial
1
0
Fork 0
mirror of https://github.com/yunjey/pytorch-tutorial.git synced 2025-07-26 03:21:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

model edited

Browse Source
This commit is contained in:
yunjey
2017-03-11 20:00:02 +09:00
parent 85d2e3e277
commit f159685a61
1 changed files with 4 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
tutorials/09 - Image Captioning/model.py
View File

@ -1,7 +1,7 @@
import torch import torch
import torch.nn as nn import torch.nn as nn
import torchvision.models as models import torchvision.models as models
import torch.nn.utils.rnn as rnn_utils from torch.nn.utils.rnn import pack_padded_sequence
from torch.autograd import Variable from torch.autograd import Variable
@ -31,27 +31,22 @@ class DecoderRNN(nn.Module):
self.lstm = nn.LSTM(embed_size, hidden_size, num_layers) self.lstm = nn.LSTM(embed_size, hidden_size, num_layers)
self.linear = nn.Linear(hidden_size, vocab_size) self.linear = nn.Linear(hidden_size, vocab_size)
def init_weights(self):
pass
def forward(self, features, captions, lengths): def forward(self, features, captions, lengths):
"""Decode image feature vectors and generate caption.""" """Decode image feature vectors and generate caption."""
embeddings = self.embed(captions) embeddings = self.embed(captions)
embeddings = torch.cat((features.unsqueeze(1), embeddings), 1) embeddings = torch.cat((features.unsqueeze(1), embeddings), 1)
packed = rnn_utils.pack_padded_sequence(embeddings, lengths, batch_first=True) # lengths is ok packed = pack_padded_sequence(embeddings, lengths, batch_first=True)
hiddens, _ = self.lstm(packed) hiddens, _ = self.lstm(packed)
outputs = self.linear(hiddens[0]) outputs = self.linear(hiddens[0])
return outputs return outputs
def sample(self, feature, state): def sample(self, feature, state):
"""Sample a caption for given a image feature.""" """Sample a caption for given a image feature."""
# (batch_size, seq_length, embed_size)
# features: (1, 128)
sampled_ids = [] sampled_ids = []
input = feature.unsqueeze(1) input = feature.unsqueeze(1)
for i in range(20): for i in range(20):
hidden, state = self.lstm(input, state) # (1, 1, 512) hidden, state = self.lstm(input, state) # (1, 1, hidden_size)
output = self.linear(hidden.view(-1, self.hidden_size)) # (1, 10000) output = self.linear(hidden.view(-1, self.hidden_size)) # (1, vocab_size)
predicted = output.max(1)[1] predicted = output.max(1)[1]
sampled_ids.append(predicted) sampled_ids.append(predicted)
input = self.embed(predicted) input = self.embed(predicted)