generate_caption.py 文件源码

def generate(net, image_model, image_path):
    feature = image_model.feature(image_path)
    net.initialize(feature)
    candidates = [(net, [bos], 0)]

    for i in range(max_length):
        next_candidates = []
        for prev_net, tokens, likelihood in candidates:
            if tokens[-1] == eos:
                next_candidates.append((None, tokens, likelihood))
                continue
            net = prev_net.copy()
            x = xp.asarray([tokens[-1]]).astype(np.int32)
            y = F.softmax(net(x))
            token_likelihood = np.log(cuda.to_cpu(y.data[0]))
            order = token_likelihood.argsort()[-beam_width:][::-1]
            next_candidates.extend([(net, tokens + [i], likelihood + token_likelihood[i]) for i in order])
        candidates = sorted(next_candidates, key=lambda x: -x[2])[:beam_width]
        if all([candidate[1][-1] == eos for candidate in candidates]):
            break
    return [candidate[1] for candidate in candidates]