batch_size = fc_feats.size(0)
state = self.init_hidden(batch_size)
seq = []
seqLogprobs = []
for t in range(self.seq_length + 2):
if t == 0:
xt = self.img_embed(fc_feats)
else:
if t == 1: // input <bos>
it = fc_feats.data.new(batch_size).long().zero_()
elif sample_max:
sampleLogprobs, it = torch.max(logprobs.data, 1)
it = it.view(-1).long()
else:
if temperature == 1.0:
prob_prev = torch.exp(logprobs.data).cpu() // fetch prev distribution: shape Nx(M+1)
else:
// scale logprobs by temperature
prob_prev = torch.exp(torch.div(logprobs.data, temperature)).cpu()
it = torch.multinomial(prob_prev, 1).cuda()
sampleLogprobs = logprobs.gather(1, it) // gather the logprobs at sampled positions
it = it.view(-1).long() // and flatten indices for downstream processing
xt = self.embed(it)
if t >= 2:
// stop when all finished
if t == 2:
unfinished = it > 0
else:
unfinished = unfinished * (it > 0)
if unfinished.sum() == 0:
break
it = it * unfinished.type_as(it)
seq.append(it) //seq[t] the input of t+2 time step
seqLogprobs.append(sampleLogprobs.view(-1))
output, state = self.core(xt.unsqueeze(0), state)
logprobs = F.log_softmax(self.logit(self.dropout(output.squeeze(0))), dim=1)
return torch.cat([_.unsqueeze(1) for _ in seq], 1), torch.cat([_.unsqueeze(1) for _ in seqLogprobs], 1)
After Change
batch_size = fc_feats.size(0)
state = self.init_hidden(batch_size)
seq = fc_feats.new_zeros(batch_size, self.seq_length, dtype=torch.long)
seqLogprobs = fc_feats.new_zeros(batch_size, self.seq_length)
for t in range(self.seq_length + 2):
if t == 0:
xt = self.img_embed(fc_feats)
else:
if t == 1: // input <bos>
it = fc_feats.data.new(batch_size).long().zero_()
xt = self.embed(it)
output, state = self.core(xt.unsqueeze(0), state)
logprobs = F.log_softmax(self.logit(self.dropout(output.squeeze(0))), dim=1)
// sample the next word
if t == self.seq_length + 1: // skip if we achieve maximum length
break
if sample_max:
sampleLogprobs, it = torch.max(logprobs.data, 1)
it = it.view(-1).long()
else:
if temperature == 1.0:
prob_prev = torch.exp(logprobs.data).cpu() // fetch prev distribution: shape Nx(M+1)
else:
// scale logprobs by temperature
prob_prev = torch.exp(torch.div(logprobs.data, temperature)).cpu()
it = torch.multinomial(prob_prev, 1).cuda()
sampleLogprobs = logprobs.gather(1, it) // gather the logprobs at sampled positions
it = it.view(-1).long() // and flatten indices for downstream processing
if t >= 1:
// stop when all finished
if t == 1:
unfinished = it > 0
else:
unfinished = unfinished * (it > 0)
it = it * unfinished.type_as(it)
seq[:,t-1] = it //seq[t] the input of t+2 time step
seqLogprobs[:,t-1] = sampleLogprobs.view(-1)
if unfinished.sum() == 0:
break
return seq, seqLogprobs
