Given x of shape (5, 400) and y of shape (40, 700) and a tensor d ([11,5,8,7,9]), I’m looking to repeat each row of x with each value in d so that the total length matches with y.
The below code works. But, can we do any better than using for loop in forward function?
import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
x = np.random.randn(5, 400)
x = np.array(x, np.float32)
xt = Variable(torch.from_numpy(x))
y = np.random.randn(40, 700)
y = np.array(y, np.float32)
yt = Variable(torch.from_numpy(y))
d = np.array([11,5,8,7,9], 'int32')
dt = torch.ByteTensor(d)
class model(nn.Module):
def __init__(self, in_dim, out_dim):
super(model, self).__init__()
self.in_dim = in_dim
self.out_dim = out_dim
self.fc1 = nn.Linear(in_dim, 256)
self.fc2 = nn.Linear(256, out_dim)
def forward(self, x, d, n):
h1 = self.fc1(x)
new_h1 = h1[0].repeat(d[0], 1)
for i in range(1, n):
new_h1 = torch.cat((new_h1, h1[i].repeat(d[i], 1)), 0)
h1 = new_h1
h2 = self.fc2(h1)
return h2
m = model(400, 700)
y_pred = m(xt, dt, len(d))
print(y_pred.size()) # 40 x 700