I would like to do a tensor split in pytorch. However, I get an error message because I can't get the splitting to work.
The behavior I want is to split the input data into two Fully Connected layers. I then want to create a model that combines the two Fully Connected layers into one. I believe the error is due to a wrong code in x1, x2 = torch.tensor_split(x,2)

import torch
from torch import nn, optim
import numpy as np
from matplotlib import pyplot as plt

class Regression(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(1, 32)
        self.linear2 = nn.Linear(32, 16)
        self.linear3 = nn.Linear(16*2, 1)

    def forward(self, x):
        x1, x2 = torch.tensor_split(x,2)
        x1 = nn.functional.relu(self.linear1(x1))
        x1 = nn.functional.relu(self.linear2(x1))

        x2 = nn.functional.relu(self.linear1(x2))
        x2 = nn.functional.relu(self.linear2(x2))
        cat_x = torch.cat([x1, x2], dim=1)

        cat_x = self.linear3(cat_x)
        return cat_x

def train(model, optimizer, E, iteration, x, y):
    losses = []
    for i in range(iteration):
        optimizer.zero_grad()                   # 勾配情報を0に初期化
        y_pred = model(x)                       # 予測
        loss = E(y_pred.reshape(y.shape), y)    # 損失を計算(shapeを揃える)
        loss.backward()                         # 勾配の計算
        optimizer.step()                        # 勾配の更新
        losses.append(loss.item())              # 損失値の蓄積
        print('epoch=', i 1, 'loss=', loss)
    return model, losses

x = np.random.uniform(0, 10, 100)                                   # x軸をランダムで作成
y = np.random.uniform(0.9, 1.1, 100) * np.sin(2 * np.pi * 0.1 * x)  # 正弦波を作成
x = torch.from_numpy(x.astype(np.float32)).float()                  # xをテンソルに変換
y = torch.from_numpy(y.astype(np.float32)).float()                  # yをテンソルに変換
X = torch.stack([torch.ones(100), x], 1)   

net = Regression()

optimizer = optim.RMSprop(net.parameters(), lr=0.01)                # 最適化にRMSpropを設定
E = nn.MSELoss()                                                    # 損失関数にMSEを設定

net, losses = train(model=net, optimizer=optimizer, E=E, iteration=5000, x=X, y=y)

error message

/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in linear(input, weight, bias)
   1846     if has_torch_function_variadic(input, weight, bias):
   1847         return handle_torch_function(linear, (input, weight, bias), input, weight, bias=bias)
-> 1848     return torch._C._nn.linear(input, weight, bias)
   1849 
   1850 

RuntimeError: mat1 and mat2 shapes cannot be multiplied (50x2 and 1x32)

CodePudding user response：

Tl;dr

Specify dim=1 in torch.tensor_split(x,2) .

Explanation

The x comes from two tensors with the shape [100,1] stacked at dim 1, so its shape is [100, 2]. After applying tensor_split, you get two tensors both with shape [50, 2].

print(x.shape) # torch.Size([100, 2])
print(torch.tensor_split(X,2)[0].shape) # torch.Size([50, 2])

The error occurred because linear1 only accepts tensors with shape [BATCH_SIZE,1] as the input, but a tensor with shape [50, 2] was passed in.

If your intention was to split the array of random numbers and the array of all ones, change torch.tensor_split(x,2) to torch.tensor_split(x,2,dim=1), which produces two tensors with the shape [100,1].