如何在Batch、PyTorch上加零

Question

有更好的方法吗？如何用零填充张量，而不创建新的张量对象？我需要输入始终是相同的batchsize，所以我想用零填充小于batchsize的输入。当序列长度较短时，就像在NLP中填充零一样，但这是对批处理的填充。

目前，我创建了一个新的张量，但正因为如此，我的GPU将耗尽内存。我不想为了处理这个操作而将批大小减少一半。

import torch
from torch import nn
class MyModel(nn.Module):
def __init__(self, batchsize=16):
super().__init__()
self.batchsize = batchsize

def forward(self, x):
b, d = x.shape

print(x.shape) # torch.Size([7, 32])
if b != self.batchsize: # 2. I need batches to be of size 16, if batch isn't 16, I want to pad the rest to zero
new_x = torch.zeros(self.batchsize,d) # 3. so I create a new tensor, but this is bad as it increase the GPU memory required greatly
new_x[0:b,:] = x
x = new_x
b = self.batchsize

print(x.shape) # torch.Size([16, 32])
return x
model = MyModel()
x = torch.randn((7, 32)) # 1. shape's batch is 7, because this is last batch, and I dont want to "drop_last"
y = model(x)
print(y.shape)

Answer 1

您可以填充额外的元素，如：

import torch.nn.functional as F
n = self.batchsize - b
new_x = F.pad(x, (0,0,n,0)) # pad the start of 2d tensors
new_x = F.pad(x, (0,0,0,n)) # pad the end of 2d tensors
new_x = F.pad(x, (0,0,0,0,0,n)) # pad the end of 3d tensors