ミニバッチ学習がうまくいかない(PyTorchでSin波回帰)

Question

前提・実現したいこと

PyTorchでSin波の回帰をしたいです．
回帰は初心者ですが，分類問題は実装をしたことがあり，
そのときのソースコードを転用(損失関数等変更)しています．

発生している問題・エラーメッセージ

バッチサイズが1のときとりあえず学習が進む(Loss<0.01)のに，
バッチサイズを2以上にすると学習が進みません(Loss>0.1)．
以下，関係あるかはわかりませんが出ている警告文です．

Using a target size (torch.Size([10])) that is different to
the input size (torch.Size([10, 1])).
This will likely lead to incorrect results due to broadcasting.
Please ensure they have the same size.
  return F.mse_loss(input, target, reduction=self.reduction)

該当のソースコード

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data as data

seed = 0

class MLPNet(nn.Module):
    def __init__(self, N_r, N_y):
        super(MLPNet, self).__init__()
        self.fc1 = nn.Linear(N_r, 100)
        self.fc2 = nn.Linear(100,100)
        self.fc3 = nn.Linear(100,100)
        self.fc4 = nn.Linear(100, N_y)
        self.dropout1 = nn.Dropout2d(0.2)
        self.dropout2 = nn.Dropout2d(0.2)
        self.dropout3 = nn.Dropout2d(0.2)
        
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.dropout1(x)
        x = F.relu(self.fc2(x))
        x = self.dropout2(x)
        x = F.relu(self.fc3(x))
        x = self.dropout3(x)        
        return self.fc4(x)
    
class MLP:
    def __init__(self, N_r, N_y, batch=10, epoch=100):
        self.N_r = N_r
        self.N_y = N_y
        self.batch = batch
        self.epoch = epoch
        
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        torch.manual_seed(seed)

        self.model = MLPNet(N_r, N_y).to(self.device)
        self.criterion = nn.MSELoss()
        self.optimizer = optim.Adam(self.model.parameters(), lr=0.01) 
    
    def train(self, R, Y):
        dataset = data.TensorDataset(torch.from_numpy(R).float(),
　　　　　　　　　　　　　　　　　　　　torch.from_numpy(Y).float().reshape(-1))
        dataloader = data.DataLoader(dataset=dataset, batch_size=self.batch, shuffle=True)        
        self.model.train()
        for e in range(self.epoch):
            for i, (r, y) in enumerate(dataloader):
                r, y = r.to(self.device), y.to(self.device)
                self.optimizer.zero_grad()
                output = self.model(r)
                loss = self.criterion(output, y)
                loss.backward()
                self.optimizer.step()
            if(e%10==0):
                print ('Epoch [%d/%d], Loss: %.4f'  % (e, self.epoch, loss.item()))

if __name__ == "__main__": 
    output = MLP(2,1)
    X = np.random.uniform(0, 5, 200)
    Y = np.sin(X).reshape(-1,1)
    X = np.hstack((X.reshape(-1,1), np.ones(200).reshape(-1,1))) #定数項を追加
    output.train(X,Y)

試したこと

dataloaderのshuffle=Falseにもしてみました．
学習率の変更もしてみました．
rとyをprintして正しくsinの関係になっていることを確認しました．
バッチサイズ2以上のときの学習後モデル出力を確認すると入力によらず~0でした．

同じソースコードで分類問題は解けています
(損失関数はCrossEntropyLoss，datasetのYについてfloatではなくlong)．

補足情報（FW/ツールのバージョンなど）

Spyder(Python3.8)
numpy 1.19.5
torch 1.9.0

Answer 1

質問者様が懸念していた警告の影響で、正常な学習データの並びと異なる並びで学習されてしまっていたようです。

以下を修正すると、警告が取れて、学習も進みますので、お試しください。

修正前

        dataset = data.TensorDataset(torch.from_numpy(R).float(),
　　　　　　　　　　　　　　　　　　　　torch.from_numpy(Y).float().reshape(-1))

修正後

        dataset = data.TensorDataset(torch.from_numpy(R).float(),
            torch.from_numpy(Y).float().reshape(-1, 1))

Comments

[ususio_ot]

正しく学習が進むことを確認いたしました．
rとyがsinの関係になっていたので無関係だろうと思い込んでしまっていました．
警告文そのままが原因で申し訳ないです．助かりました，ありがとうございました．