Google Colab上で実行したサンプルコードです。 多層パーセプトロンで回帰問題を解きます。 データセットは単純な数式に基づくものを、ランダム生成しています。 ここの複雑度をいじることで、DNNの表現力を測ることを目的にしています。

import torch.nn.functional as F
import torch
import random
from torch.utils.data import DataLoader
 
# 独自のデータセットクラス
class MyDataset(torch.utils.data.Dataset):
    def __init__(self, param_a_list, param_b_list, param_c_list, ans_list):
        self._param_a = param_a_list
        self._param_b = param_b_list
        self._param_c = param_c_list
        self._ans = ans_list
 
    def __len__(self):
        return len(self._ans)
 
    def __getitem__(self, idx):
        input_tensor = torch.FloatTensor([self._param_a[idx], self._param_b[idx], self._param_c[idx]])
        ans_tensor = torch.FloatTensor([self._ans[idx]])
        return input_tensor, ans_tensor
 
# DNNモデル(多層パーセプトロン)
class MLPR(torch.nn.Module):
    def __init__(self, n_input, n_hidden, n_output):
        super(MLPR, self).__init__()
        self.l1 = torch.nn.Linear(n_input, n_hidden)
        self.l2 = torch.nn.Linear(n_hidden, n_hidden)
        self.l3 = torch.nn.Linear(n_hidden, n_output)
 
    def forward(self, x):
        h1 = self.l1(x)
        h2 = torch.relu(h1)
        h3 = self.l2(h2)
        h4 = torch.relu(h3)
        h5 = self.l3(h4)
        return h5
 
# 学習データを作る
param_a_list = []
param_b_list = []
param_c_list = []
ans_list = []
 
for i in range(0, 2000):
    prm_a = random.random()
    prm_b = random.random()
    prm_c = random.random()
 
    # ここから学習させたいロジック部分
    ans = prm_a * 5 + prm_b / 4 - prm_c * 2
    if ans > 0.5:
      ans = ans * 2
    # ここまで
 
    param_a_list.append(prm_a)
    param_b_list.append(prm_b)
    param_c_list.append(prm_c)
    ans_list.append(ans)
 
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('using device:', device)
 
myds = MyDataset(param_a_list, param_b_list, param_c_list, ans_list)
train_loader = DataLoader(myds, batch_size=100, shuffle=True)
 
# 入力は3、出力は1
model = MLPR(3, 500, 1)
# 回帰問題なのでとりあえずMSELossを損失関数に使う
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
 
model.to(device)
criterion.to(device)
 
# 学習プロセス
# 学習が進みに応じた損失関数の変化を見る。検証プロセスは無し。
from torch.autograd import Variable
loss_history = []
for epoch in range(2000):
    total_loss = 0
    for x_train, y_train in train_loader:
        # Variableは微分機能内包
        x_train = Variable(x_train)
        y_train = Variable(y_train)
        optimizer.zero_grad()
        x_train = x_train.to(device)
        y_train = y_train.to(device)
        y_pred = model(x_train)
        loss = criterion(y_pred, y_train)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    loss_history.append(total_loss)
    if (epoch +1) % 100 == 0:
        print(epoch + 1, total_loss)