PyTorch入門 第10回：GANとRNNをまとめて理解する

第9回：画像タスクの出力を読む
  -> logits、mask、box、score
  -> argmax、NMS、IoU
  -> train/evalで戻り値が変わるモデル

第10回：生成と系列の計算グラフを読む
    -> fake.detach で勾配の流れを止める
  -> GeneratorとDiscriminatorを別々に更新する
  -> hidden stateを時間方向に受け渡す
  -> paddingとpackingで系列長の違いを扱う

z ~ p_z(z)

x_fake = G(z)

D(x) -> real/fake logit

latent z
 -> Linear / ConvTranspose / Upsample blocks
 -> generated sample G(z)

real x or fake G(z)
 -> feature extractor
 -> scalar logit

min_G max_D V(D, G)
= E_x~p_data [log D(x)]
  + E_z~p_z [log(1 - D(G(z)))]

D*(x) = p_data(x) / (p_data(x) + p_g(x))

p_g = p_data
D(x) = 1/2

real x -> D -> target 1
noise z -> G(z).detach() -> D -> target 0
loss_D = loss_real + loss_fake

optimizer_d.zero_grad(set_to_none=True)

real_logits = discriminator(real_images)
real_loss = criterion(real_logits, torch.ones_like(real_logits))

noise = torch.randn(batch_size, latent_dim, device=device)
fake_samples = generator(noise)
fake_logits = discriminator(fake_samples.detach())  # D更新ではGへ勾配を戻さない
fake_loss = criterion(fake_logits, torch.zeros_like(fake_logits))

loss_d = real_loss + fake_loss
loss_d.backward()
optimizer_d.step()

noise z -> G(z) -> D -> target 1
loss_G = BCEWithLogitsLoss(D(G(z)), ones)

optimizer_g.zero_grad(set_to_none=True)

fake_logits_for_g = discriminator(fake_samples)
loss_g = criterion(fake_logits_for_g, torch.ones_like(fake_logits_for_g))
loss_g.backward()
optimizer_g.step()

import torch
from torch import nn


logits = torch.tensor([[2.0], [-1.0], [0.0]])
targets = torch.tensor([[1.0], [0.0], [1.0]])

criterion = nn.BCEWithLogitsLoss()
loss = criterion(logits, targets)

print("loss:", round(loss.item(), 4))
# loss: 0.3778

推奨しやすい形:
  discriminator -> logits
  BCEWithLogitsLoss(logits, target)

避けたい形:
  discriminator -> sigmoid probability
  BCELoss(probability, target)

import torch
from torch import nn


torch.manual_seed(0)

latent_dim = 3
data_dim = 4
batch_size = 5

generator = nn.Sequential(nn.Linear(latent_dim, data_dim))
discriminator = nn.Sequential(nn.Linear(data_dim, 1))
criterion = nn.BCEWithLogitsLoss()
optimizer_g = torch.optim.SGD(generator.parameters(), lr=0.1)
optimizer_d = torch.optim.SGD(discriminator.parameters(), lr=0.1)

real_samples = torch.randn(batch_size, data_dim)
noise = torch.randn(batch_size, latent_dim)
fake_samples = generator(noise)

# 1. Discriminatorを更新する
optimizer_d.zero_grad(set_to_none=True)
real_logits = discriminator(real_samples)
fake_logits = discriminator(fake_samples.detach())  # このstepではGへ勾配を流さない
loss_d = criterion(real_logits, torch.ones_like(real_logits)) + criterion(
    fake_logits,
    torch.zeros_like(fake_logits),
)
loss_d.backward()
optimizer_d.step()

print("G grad after D step:", generator[0].weight.grad is None)
# G grad after D step: True

# 2. Generatorを更新する
optimizer_g.zero_grad(set_to_none=True)
fake_logits_for_g = discriminator(fake_samples)  # ここではdetachしない
loss_g = criterion(fake_logits_for_g, torch.ones_like(fake_logits_for_g))
loss_g.backward()
optimizer_g.step()

print("G grad after G step:", generator[0].weight.grad.abs().sum().item() > 0)
# G grad after G step: True

z: (N, nz, 1, 1)
 -> ConvTranspose2d + BatchNorm + ReLU
 -> ConvTranspose2d + BatchNorm + ReLU
 -> ConvTranspose2d + BatchNorm + ReLU
 -> ConvTranspose2d + Tanh
 -> image: (N, nc, 64, 64)

image: (N, nc, 64, 64)
 -> strided Conv2d + LeakyReLU
 -> strided Conv2d + BatchNorm + LeakyReLU
 -> strided Conv2d + BatchNorm + LeakyReLU
 -> Conv2d
 -> scalar logit

import torch
from torch import nn


class MiniGenerator(nn.Module):
    def __init__(self, nz=16, ngf=8, nc=3):
        super().__init__()
        self.main = nn.Sequential(
            nn.ConvTranspose2d(nz, ngf * 8, kernel_size=4, stride=1, padding=0, bias=False),  # (N, nz, 1, 1) -> (N, ngf*8, 4, 4)
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(ngf * 8, ngf * 4, kernel_size=4, stride=2, padding=1, bias=False),  # 4x4 -> 8x8
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(ngf * 4, ngf * 2, kernel_size=4, stride=2, padding=1, bias=False),  # 8x8 -> 16x16
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(ngf * 2, ngf, kernel_size=4, stride=2, padding=1, bias=False),  # 16x16 -> 32x32
            nn.BatchNorm2d(ngf),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(ngf, nc, kernel_size=4, stride=2, padding=1, bias=False),  # 32x32 -> 64x64
            nn.Tanh(),
        )

    def forward(self, z):
        return self.main(z)


class MiniDiscriminator(nn.Module):
    def __init__(self, nc=3, ndf=8):
        super().__init__()
        self.main = nn.Sequential(
            nn.Conv2d(nc, ndf, kernel_size=4, stride=2, padding=1, bias=False),  # 64x64 -> 32x32
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf, ndf * 2, kernel_size=4, stride=2, padding=1, bias=False),  # 32x32 -> 16x16
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 2, ndf * 4, kernel_size=4, stride=2, padding=1, bias=False),  # 16x16 -> 8x8
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 4, ndf * 8, kernel_size=4, stride=2, padding=1, bias=False),  # 8x8 -> 4x4
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 8, 1, kernel_size=4, stride=1, padding=0, bias=False),  # 4x4 -> 1x1
        )

    def forward(self, image):
        return self.main(image)


z = torch.randn(2, 16, 1, 1)
generator = MiniGenerator()
discriminator = MiniDiscriminator()

fake_images = generator(z)
logits = discriminator(fake_images)

print("fake images:", tuple(fake_images.shape), fake_images.min().item() >= -1.0, fake_images.max().item() <= 1.0)
# fake images: (2, 3, 64, 64) True True
print("logits:", tuple(logits.shape))
# logits: (2, 1, 1, 1)

import torch


values = torch.tensor([0.0, 0.5, 1.0])
scaled = (values - 0.5) / 0.5

print(scaled.tolist())
# [-1.0, 0.0, 1.0]

uint8 0..255
 -> float 0..1
 -> Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
 -> float -1..1

from torch import nn


def weights_init(module):
    classname = module.__class__.__name__
    if "Conv" in classname:
        nn.init.normal_(module.weight, mean=0.0, std=0.02)
    elif "BatchNorm" in classname:
        nn.init.normal_(module.weight, mean=1.0, std=0.02)
        nn.init.constant_(module.bias, 0.0)

/path/to/celeba
  -> img_align_celeba
       -> 188242.jpg
       -> 173822.jpg
       -> ...

from pathlib import Path
import imageio.v2 as imageio
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.utils import make_grid

from tools.dcgan import Discriminator, Generator
from tools.my_dataset import CelebADataset


torch.manual_seed(1)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
data_dir = Path("data/celeba/img_align_celeba")  # 画像ファイルを置いたディレクトリに合わせる
out_dir = Path("log_gan")
out_dir.mkdir(parents=True, exist_ok=True)

image_size = 64
nc = 3
nz = 100
ngf = 128
ndf = 128
num_epochs = 20
batch_size = 64
lr = 0.0002
beta1 = 0.5
checkpoint_interval = 10

fixed_noise = torch.randn(64, nz, 1, 1, device=device)  # 各epochで同じzを使い、生成の変化を比べる

transform = transforms.Compose([
    transforms.Resize(image_size),
    transforms.CenterCrop(image_size),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # 画像値域を[-1, 1]へそろえる
])

train_set = CelebADataset(data_dir=data_dir, transforms=transform)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2)

net_g = Generator(nz=nz, ngf=ngf, nc=nc).to(device)
net_d = Discriminator(nc=nc, ndf=ndf).to(device)
net_g.initialize_weights()
net_d.initialize_weights()

criterion = nn.BCELoss()
optimizer_g = torch.optim.Adam(net_g.parameters(), lr=lr, betas=(beta1, 0.999))
optimizer_d = torch.optim.Adam(net_d.parameters(), lr=lr, betas=(beta1, 0.999))
scheduler_g = torch.optim.lr_scheduler.StepLR(optimizer_g, step_size=8, gamma=0.1)
scheduler_d = torch.optim.lr_scheduler.StepLR(optimizer_d, step_size=8, gamma=0.1)

g_losses = []
d_losses = []

for epoch in range(num_epochs):
    net_g.train()
    net_d.train()

    for batch_index, real_images in enumerate(train_loader):
        real_images = real_images.to(device)
        batch_current = real_images.size(0)
        real_labels = torch.ones(batch_current, device=device)
        fake_labels = torch.zeros(batch_current, device=device)

        # 1. Discriminatorを更新する: realは1、fakeは0へ近づける
        optimizer_d.zero_grad(set_to_none=True)
        noise = torch.randn(batch_current, nz, 1, 1, device=device)
        fake_images = net_g(noise)

        real_probs = net_d(real_images).view(-1)
        fake_probs = net_d(fake_images.detach()).view(-1)  # D更新ではGへ勾配を戻さない
        loss_d_real = criterion(real_probs, real_labels)
        loss_d_fake = criterion(fake_probs, fake_labels)
        loss_d = loss_d_real + loss_d_fake
        loss_d.backward()
        optimizer_d.step()

        d_x = real_probs.mean().item()
        d_g_z_before = fake_probs.mean().item()

        # 2. Generatorを更新する: fakeをrealと判定させる方向へGを動かす
        optimizer_g.zero_grad(set_to_none=True)
        fake_probs_for_g = net_d(fake_images).view(-1)  # G更新ではdetachしない
        loss_g = criterion(fake_probs_for_g, real_labels)
        loss_g.backward()
        optimizer_g.step()

        d_g_z_after = fake_probs_for_g.mean().item()

        if batch_index % 10 == 0:
            print(
                f"[{epoch}/{num_epochs}][{batch_index}/{len(train_loader)}]\t"
                f"Loss_D: {loss_d.item():.4f}\tLoss_G: {loss_g.item():.4f}\t"
                f"D(x): {d_x:.4f}\tD(G(z)): {d_g_z_before:.4f} / {d_g_z_after:.4f}"
            )

        g_losses.append(loss_g.item())
        d_losses.append(loss_d.item())

    scheduler_d.step()
    scheduler_g.step()

    # 固定noiseから生成した画像を保存し、epochごとの変化を見る
    net_g.eval()
    with torch.inference_mode():
        fake = net_g(fixed_noise).detach().cpu()
    image_grid = make_grid(fake, padding=2, normalize=True).numpy()
    image_grid = np.transpose(image_grid, (1, 2, 0))
    plt.imshow(image_grid)
    plt.axis("off")
    plt.title(f"Epoch: {epoch}")
    plt.savefig(out_dir / f"{epoch}_epoch.png")
    plt.close()

    if (epoch + 1) % checkpoint_interval == 0:
        checkpoint = {
            "g_model_state_dict": net_g.state_dict(),
            "d_model_state_dict": net_d.state_dict(),
            "epoch": epoch,
        }
        torch.save(checkpoint, out_dir / f"checkpoint_{epoch}_epoch.pkl")

plt.figure(figsize=(10, 5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(g_losses, label="G")
plt.plot(d_losses, label="D")
plt.xlabel("iterations")
plt.ylabel("Loss")
plt.legend()
plt.savefig(out_dir / "loss.png")
plt.close()

epoch_images = sorted(out_dir.glob("*_epoch.png"), key=lambda path: int(path.stem.split("_")[0]))
frames = [imageio.imread(path) for path in epoch_images]
imageio.mimsave(out_dir / "generation_animation.gif", frames, fps=2)

[0/20][0/32]    Loss_D: 1.3862    Loss_G: 0.7421    D(x): 0.5010    D(G(z)): 0.5004 / 0.4762

from collections import OrderedDict
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import torch
from torchvision.utils import make_grid

from tools.dcgan import Generator


def remove_module_prefix(state_dict):
    cleaned_state = OrderedDict()
    for key, value in state_dict.items():
        cleaned_key = key.removeprefix("module.")
        cleaned_state[cleaned_key] = value
    return cleaned_state


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
path_checkpoint = Path("log_gan/checkpoint_19_epoch.pkl")  # 実際に保存したcheckpoint名に合わせる

num_images = 64
nz = 100
ngf = 128
nc = 3

net_g = Generator(nz=nz, ngf=ngf, nc=nc)
checkpoint = torch.load(path_checkpoint, map_location="cpu", weights_only=True)
state_dict_g = remove_module_prefix(checkpoint["g_model_state_dict"])
net_g.load_state_dict(state_dict_g)
net_g.to(device).eval()

fixed_noise = torch.randn(num_images, nz, 1, 1, device=device)

with torch.inference_mode():
    fake_images = net_g(fixed_noise).detach().cpu()

image_grid = make_grid(fake_images, padding=2, normalize=True).numpy()
image_grid = np.transpose(image_grid, (1, 2, 0))
plt.imshow(image_grid)
plt.axis("off")
plt.show()

many z -> nearly identical output

x_t + h_(t-1) -> recurrent cell -> h_t -> optional output y_t

h_t = tanh(W_ih x_t + b_ih + W_hh h_(t-1) + b_hh)
y_t = W_ho h_t + b_ho

x_t: (1, 58)
h_t: (1, 128)
y_t: (1, 18)

h0 -> [cell x1] -> h1 -> [cell x2] -> h2 -> [cell x3] -> h3
          same W             same W             same W

1. 系列を前から読む
2. 各時刻のhidden stateを作る
3. 最後または各時刻の出力からlossを計算する
4. 時間方向にさかのぼって勾配を流す
5. 共有パラメータに勾配が合算される

hidden = hidden.detach()

import torch
from torch import nn


torch.manual_seed(0)

rnn = nn.RNN(
    input_size=5,
    hidden_size=7,
    num_layers=2,  # h_nの先頭次元 D * num_layers に影響する
    batch_first=True,
    bidirectional=True,  # 双方向なのでD=2
)

x = torch.randn(3, 4, 5)  # (N, L, H_in)
output, h_n = rnn(x)

print("output:", tuple(output.shape))
# output: (3, 4, 14)
print("h_n:", tuple(h_n.shape))
# h_n: (4, 3, 7)

input : (N, L, H_in)
output: (N, L, D * H_out)
h_n   : (D * num_layers, N, H_out)

i_t = sigmoid(input gate)
f_t = sigmoid(forget gate)
g_t = tanh(candidate)
o_t = sigmoid(output gate)

c_t = f_t * c_(t-1) + i_t * g_t
h_t = o_t * tanh(c_t)

import torch
from torch import nn


torch.manual_seed(0)

lstm = nn.LSTM(
    input_size=5,
    hidden_size=7,
    num_layers=1,
    batch_first=True,
)

x = torch.randn(3, 4, 5)
output, (h_n, c_n) = lstm(x)

print("output:", tuple(output.shape))
# output: (3, 4, 7)
print("h_n:", tuple(h_n.shape), "c_n:", tuple(c_n.shape))
# h_n: (1, 3, 7) c_n: (1, 3, 7)

r_t = sigmoid(...)
z_t = sigmoid(...)
n_t = tanh(... + r_t * (...))
h_t = (1 - z_t) * n_t + z_t * h_(t-1)

import torch
from torch import nn


torch.manual_seed(0)

gru = nn.GRU(
    input_size=5,
    hidden_size=7,
    num_layers=1,
    batch_first=True,
)

x = torch.randn(3, 4, 5)
output, h_n = gru(x)

print("output:", tuple(output.shape))
# output: (3, 4, 7)
print("h_n:", tuple(h_n.shape))
# h_n: (1, 3, 7)

Chou      length 4
Smith     length 5
Li        length 2

padded -> shape (N, L_max, feature)

import torch
from torch import nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence


torch.manual_seed(0)

token_ids = torch.tensor(
    [[1, 2, 3, 4, 0],
     [5, 6, 7, 0, 0],
     [8, 9, 0, 0, 0]]
)
lengths = torch.tensor([4, 3, 2])

embedding = nn.Embedding(num_embeddings=10, embedding_dim=4, padding_idx=0)  # padding_idx=0の行はEmbeddingの勾配更新から除外される
gru = nn.GRU(input_size=4, hidden_size=6, batch_first=True)

embedded = embedding(token_ids)
packed = pack_padded_sequence(
    embedded,
    lengths.cpu(),  # lengthsがTensorの場合はCPU上に置く
    batch_first=True,
    enforce_sorted=False,  # 長さの降順に並べ替え済みでなくてもよい
)
packed_output, h_n = gru(packed)  # PackedSequenceを渡すとpadding部分の計算をスキップする
output, output_lengths = pad_packed_sequence(packed_output, batch_first=True)

print("embedded:", tuple(embedded.shape))
# embedded: (3, 5, 4)
print("output:", tuple(output.shape), output_lengths.tolist())
# output: (3, 4, 6) [4, 3, 2]
print("h_n:", tuple(h_n.shape))
# h_n: (1, 3, 6)

character -> integer index -> one-hot vector

import string
import unicodedata
import torch


allowed_characters = string.ascii_letters + " .,;'" + "_"
n_letters = len(allowed_characters)


def unicode_to_ascii(text):
    return "".join(
        character
        for character in unicodedata.normalize("NFD", text)
        if unicodedata.category(character) != "Mn"
        and character in allowed_characters
    )


def letter_to_index(letter):
    if letter not in allowed_characters:
        return allowed_characters.find("_")
    return allowed_characters.find(letter)


def line_to_tensor(line):
    tensor = torch.zeros(len(line), 1, n_letters)
    for position, letter in enumerate(line):
        tensor[position, 0, letter_to_index(letter)] = 1.0
    return tensor


print("n_letters:", n_letters)
# n_letters: 58
print(unicode_to_ascii("Ślusàrski"))
# Slusarski
name_tensor = line_to_tensor("Ahn")
print("shape:", tuple(name_tensor.shape))
# shape: (3, 1, 58)
print("active indices:", name_tensor.argmax(dim=2).squeeze(1).tolist())
# active indices: [26, 7, 13]

string
 -> Unicode normalization
 -> character/token IDs
 -> padded ID batch
 -> Embedding
 -> RNN/LSTM/GRU
 -> sequence representation
 -> Linear classifier

"Chou"
 -> C, h, o, u
 -> 各文字をIDまたはone-hotへ変換
 -> RNNへ順に入力
 -> 最終hiddenまたは最終出力
 -> 18クラスのlogits
 -> language label

h0
 -> cell(C, h0) -> h1
 -> cell(h, h1) -> h2
 -> cell(o, h2) -> h3
 -> cell(u, h3) -> h4
 -> classifier(h4)

from pathlib import Path
import math
import random
import string
import time
import unicodedata

import torch
from torch import nn


torch.manual_seed(1)
random.seed(1)
device = torch.device("cpu")

data_dir = Path("data/names")  # data.zipを展開したdata/namesに合わせる
all_letters = string.ascii_letters + " .,;'"
n_letters = len(all_letters)  # 57文字（アルファベット52+空白と.,;'の5文字。ここでは未知文字用の記号は含めない）


def unicode_to_ascii(text):
    return "".join(
        character
        for character in unicodedata.normalize("NFD", text)
        if unicodedata.category(character) != "Mn"
        and character in all_letters
    )


def read_lines(path):
    lines = path.read_text(encoding="utf-8").strip().split("\n")
    return [unicode_to_ascii(line) for line in lines]


def letter_to_index(letter):
    return all_letters.find(letter)


def line_to_tensor(line):
    tensor = torch.zeros(len(line), 1, n_letters)
    for position, letter in enumerate(line):
        tensor[position, 0, letter_to_index(letter)] = 1.0
    return tensor


def time_since(start_time):
    elapsed = time.time() - start_time
    minutes = math.floor(elapsed / 60)
    seconds = elapsed - minutes * 60
    return f"{minutes}m {seconds:.0f}s"


class ManualRNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.hidden_size = hidden_size
        self.u = nn.Linear(input_size, hidden_size)  # 前掲の数式のW_ih, b_ihに相当
        self.w = nn.Linear(hidden_size, hidden_size)  # 前掲の数式のW_hh, b_hhに相当
        self.v = nn.Linear(hidden_size, output_size)  # 前掲の数式のW_ho, b_hoに相当
        self.tanh = nn.Tanh()
        self.log_softmax = nn.LogSoftmax(dim=1)

    def forward(self, inputs, hidden):
        input_to_hidden = self.u(inputs)
        hidden_to_hidden = self.w(hidden)
        hidden = self.tanh(input_to_hidden + hidden_to_hidden)  # h_t = tanh(W_ih x_t + b_ih + W_hh h_(t-1) + b_hh)
        output = self.log_softmax(self.v(hidden))  # y_t = W_ho h_t + b_hoをlog probabilityに変換
        return output, hidden

    def init_hidden(self):
        return torch.zeros(1, self.hidden_size)  # h_0はゼロベクトルから始める


category_lines = {}
for path in sorted(data_dir.glob("*.txt")):
    category = path.stem
    category_lines[category] = read_lines(path)

all_categories = sorted(category_lines)
n_categories = len(all_categories)


def random_training_example():
    category = random.choice(all_categories)
    line = random.choice(category_lines[category])
    category_tensor = torch.tensor([all_categories.index(category)], dtype=torch.long)
    line_tensor = line_to_tensor(line)
    return category, line, category_tensor, line_tensor


def category_from_output(output):
    _, top_index = output.topk(1)  # LogSoftmax出力の中で最大のlog probabilityを持つクラスを1つだけ取り出す
    category_index = top_index[0].item()
    return all_categories[category_index], category_index


def train_one_example(category_tensor, line_tensor):
    hidden = model.init_hidden().to(device)
    line_tensor = line_tensor.to(device)
    category_tensor = category_tensor.to(device)

    optimizer.zero_grad(set_to_none=True)
    for position in range(line_tensor.size(0)):
        output, hidden = model(line_tensor[position], hidden)  # 1文字ずつhiddenを更新し、ループを抜けた後はoutputに最後の文字の結果だけが残る

    loss = criterion(output, category_tensor)  # 最後の文字の出力だけをNLLLossに渡す
    loss.backward()
    optimizer.step()
    return output, loss.item()


n_hidden = 128
learning_rate = 0.005
n_iters = 200000
print_every = 5000
plot_every = 5000

model = ManualRNN(n_letters, n_hidden, n_categories).to(device)
criterion = nn.NLLLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

current_loss = 0.0
loss_history = []
start_time = time.time()

for iteration in range(1, n_iters + 1):
    category, line, category_tensor, line_tensor = random_training_example()
    output, loss = train_one_example(category_tensor, line_tensor)
    current_loss += loss

    if iteration % print_every == 0:
        guess, _ = category_from_output(output)
        correct = "yes" if guess == category else f"no ({category})"
        print(
            f"Iter: {iteration:<7} time: {time_since(start_time):>8s} "
            f"loss: {loss:.4f} name: {line:>10s} pred: {guess:>8s} label: {correct:>8s}"
        )

    if iteration % plot_every == 0:
        loss_history.append(current_loss / plot_every)
        current_loss = 0.0

torch.save(model.state_dict(), "rnn_state_dict.pkl")

Iter: 5000    time: 0m 12s loss: 2.1034 name:      Satoh pred: Japanese label: yes

def evaluate(line_tensor):
    hidden = model.init_hidden().to(device)
    line_tensor = line_tensor.to(device)

    with torch.inference_mode():
        for position in range(line_tensor.size(0)):
            output, hidden = model(line_tensor[position], hidden)
    return output


def predict(input_line, n_predictions=3):
    print(f"\n> {input_line}")
    output = evaluate(line_to_tensor(input_line))
    top_values, top_indices = output.topk(n_predictions, dim=1)

    for rank in range(n_predictions):
        value = top_values[0, rank].item()
        category_index = top_indices[0, rank].item()
        print(f"({value:.2f}) {all_categories[category_index]}")


predict("Yue Tingsong")
predict("test your name")

> Yue Tingsong
(-0.31) Chinese
(-1.84) Vietnamese
(-2.45) Korean

import torch
from torch import nn
from torch.nn.utils.rnn import pack_padded_sequence


class NameClassifier(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_size, num_classes, padding_idx):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=padding_idx)
        self.gru = nn.GRU(
            embedding_dim,
            hidden_size,
            batch_first=True,
            bidirectional=True,
        )
        self.classifier = nn.Linear(hidden_size * 2, num_classes)  # 双方向なのでhidden_size * 2

    def forward(self, token_ids, lengths):
        embedded = self.embedding(token_ids)  # (N, L) -> (N, L, embedding_dim)
        packed = pack_padded_sequence(
            embedded,
            lengths.cpu(),
            batch_first=True,
            enforce_sorted=False,
        )
        _, hidden = self.gru(packed)
        representation = torch.cat([hidden[-2], hidden[-1]], dim=1)  # 最後の層のforward/backwardを連結する
        return self.classifier(representation)


model = NameClassifier(
    vocab_size=20,
    embedding_dim=6,
    hidden_size=8,
    num_classes=4,
    padding_idx=0,
)

token_ids = torch.tensor(
    [[1, 2, 3, 0],
     [4, 5, 0, 0]]
)
lengths = torch.tensor([3, 2])
logits = model(token_ids, lengths)

print("logits:", tuple(logits.shape))
# logits: (2, 4)

import torch
from torch import nn


torch.manual_seed(0)

model = nn.Linear(3, 1)
x = torch.randn(4, 3)
y = torch.randn(4, 1)
loss = nn.MSELoss()(model(x), y)
loss.backward()

total_norm = nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # 全パラメータの勾配をまとめて1つのnormにし、max_normを超えた分だけスケールダウンする

print("clipped from:", round(float(total_norm), 4))
# clipped from: 2.5123

one-layer bidirectional:
  hidden[-2] -> forward final
  hidden[-1] -> backward final

z -> Generator -> fake sample --+
                                |
real sample --------------------+-> Discriminator -> real/fake logit

Step D: real target=1 + fake.detach target=0 -> update D
Step G: noise -> G -> D target=1 -> update G

latent (N, nz, 1, 1)
 -> ConvTranspose blocks
 -> image (N, 3, 64, 64)

image (N, 3, 64, 64)
 -> strided Conv blocks
 -> logit (N, 1, 1, 1)

x_t ----+
        +-> recurrent cell -> h_t -> output
h_t-1 --+

x1 -> [RNN] -> h1
               |
x2 -> [RNN] -> h2
               |
x3 -> [RNN] -> h3 -> classifier

characters
 -> IDs / one-hot
 -> Embedding or one-hot Tensor
 -> RNN / LSTM / GRU
 -> sequence representation
 -> class logits