pyTorch上的TimeDistributed

Keras有个TimeDistributed包装器,pytorch上用nn.Linear就能实现。老是忘在这里记录下:

给定输入in[batch, steps, in_dims],希望在每个step内Dense,然后输出out[batch, steps, out_dims],

只需要直接指定nn.Linear(in_dims, out_dims)就好了,例如:

batchs=2
steps=3
in_dims=4
out_dims=2

m = nn.Linear(in_dims, out_dims, False)
print(m.weight)
Out[16]:
Parameter containing:
tensor([[ 0.4397,  0.0982, -0.0458, -0.0480],
        [-0.1751, -0.2792, -0.2744,  0.1664]], requires_grad=True)

input = torch.ones(batchs, steps, in_dims)
m(input)
Out[17]:
tensor([[[ 0.4440, -0.5623],
         [ 0.4440, -0.5623],
         [ 0.4440, -0.5623]],
        [[ 0.4440, -0.5623],
         [ 0.4440, -0.5623],
         [ 0.4440, -0.5623]]], grad_fn=)

输出维度是[batch,step,out_dims]每个step内作dense

你可能感兴趣的:(pyTorch上的TimeDistributed)