サクッと試せるPyTorchニューラルネットワークの最小コード（Minimal Code）

「とにかくコピペだけでサクッと動かせるニューラルネットワークのコードが欲しい」というときがあるのだけど、検索してもなかなか出てきません。そこで、とにかくコピペだけ試せる最小コードを作りました。PyTorch版です。

# ライブラリのインポート
import torch
import torch.nn as  nn
import torch.optim as optim

# 入力データ
NUMBER_OF_DATA  = 64
INPUT_FEATURES  = 1000
HIDDEN_FEATURES = 100
OUTPUT_FEATURES = 10
x = torch.randn(NUMBER_OF_DATA, INPUT_FEATURES)
y = torch.randn(NUMBER_OF_DATA, OUTPUT_FEATURES)
#x.size() # torch.Size([64, 1000])
#y.size() # torch.Size([64, 10])

# モデルの定義
model = nn.Sequential(
            nn.Linear(INPUT_FEATURES, HIDDEN_FEATURES),
            nn.ReLU(),
            nn.Linear(HIDDEN_FEATURES, OUTPUT_FEATURES))

# トレーニング
LEARNING_RATE = 0.03   # 学習率： 0.03
EPOCHS = 100           # エポック数： 100

criterion = nn.MSELoss()

optimizer = optim.SGD(
    model.parameters(),
    lr=LEARNING_RATE)

def train_step(train_X, train_y):
    model.train()
    pred_y = model(train_X)
    optimizer.zero_grad()
    loss = criterion(pred_y, train_y)
    loss.backward()
    optimizer.step()
    return loss.item()

for i in range(EPOCHS):
    loss_result = train_step(x, y)
    print(i, ':', loss_result)

print('Finished Training')
#print(model.state_dict())

ちなみに試すときは、Colabスクラッチパッドが便利ですよ。