対照学習(Contrastive Learning)の理論と実装

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.font_manager as _fm
for _c in ['Hiragino Sans','Yu Gothic','Noto Sans CJK JP','IPAexGothic','Meiryo']:
    if any(_c==_f.name for _f in _fm.fontManager.ttflist):
        plt.rcParams['font.family']=_c; break
plt.rcParams['axes.unicode_minus']=False

class SimpleAugmentation:
    """シンプルなデータ拡張（数値データ用）"""
    def __init__(self, noise_scale=0.1, mask_prob=0.1):
        self.noise_scale = noise_scale
        self.mask_prob = mask_prob

    def __call__(self, x):
        # ガウスノイズの追加
        noise = torch.randn_like(x) * self.noise_scale
        x_aug = x + noise
        # ランダムマスキング
        mask = (torch.rand_like(x) > self.mask_prob).float()
        x_aug = x_aug * mask
        return x_aug

class ContrastiveDataset:
    """対照学習用データセット"""
    def __init__(self, data, augmentation):
        self.data = data
        self.augmentation = augmentation

    def __len__(self):
        return len(self.data)

    def get_pair(self, idx):
        """同じデータから2つの拡張ビューを生成"""
        x = self.data[idx]
        view1 = self.augmentation(x)
        view2 = self.augmentation(x)
        return view1, view2

    def get_batch(self, batch_size):
        """バッチを取得"""
        indices = np.random.choice(len(self.data), batch_size, replace=False)
        views1, views2 = [], []
        for idx in indices:
            v1, v2 = self.get_pair(idx)
            views1.append(v1)
            views2.append(v2)
        return torch.stack(views1), torch.stack(views2)

# 合成データの生成: 5クラスのクラスタ
np.random.seed(42)
num_classes = 5
samples_per_class = 200
dim = 20

data_list = []
labels_list = []
for c in range(num_classes):
    center = np.random.randn(dim) * 1.5
    samples = center + np.random.randn(samples_per_class, dim) * 2.0
    data_list.append(samples)
    labels_list.append(np.full(samples_per_class, c))

data = torch.FloatTensor(np.concatenate(data_list))
labels = torch.LongTensor(np.concatenate(labels_list))

# シャッフル
perm = torch.randperm(len(data))
data = data[perm]
labels = labels[perm]

print(f"データサイズ: {data.shape}")
print(f"クラス数: {num_classes}")
print(f"クラスごとのサンプル数: {samples_per_class}")

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimCLR(nn.Module):
    """SimCLRモデル"""
    def __init__(self, input_dim, hidden_dim=64, proj_dim=32):
        super().__init__()
        # エンコーダ f: 入力 -> 表現 h
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )
        # 射影ヘッド g: 表現 h -> 射影 z
        self.projector = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, proj_dim)
        )

    def forward(self, x):
        h = self.encoder(x)    # 表現（下流タスクで使用）
        z = self.projector(h)   # 射影（対照学習の損失計算で使用）
        return h, z

    def get_representation(self, x):
        """下流タスク用: エンコーダのみの出力"""
        return self.encoder(x)

def info_nce_loss(z1, z2, temperature=0.5):
    """
    InfoNCE損失の計算
    z1, z2: 正例ペアの射影表現 (batch_size, proj_dim)
    temperature: 温度パラメータ
    """
    batch_size = z1.size(0)

    # L2正規化（コサイン類似度の計算のため）
    z1 = F.normalize(z1, dim=1)
    z2 = F.normalize(z2, dim=1)

    # 全ペアの類似度行列を計算
    # z: (2*batch_size, proj_dim) に連結
    z = torch.cat([z1, z2], dim=0)  # (2N, d)
    sim_matrix = torch.matmul(z, z.t()) / temperature  # (2N, 2N)

    # 自分自身との類似度をマスク（-infにして softmax で0にする）
    mask = torch.eye(2 * batch_size, device=z.device).bool()
    sim_matrix.masked_fill_(mask, float('-inf'))

    # 正例のインデックスを設定
    # z1[i] の正例は z2[i]（インデックス i+N）
    # z2[i] の正例は z1[i]（インデックス i）
    pos_indices = torch.cat([
        torch.arange(batch_size, 2 * batch_size),  # z1の正例はz2
        torch.arange(0, batch_size)                 # z2の正例はz1
    ]).to(z.device)

    # InfoNCE損失 = クロスエントロピー損失
    loss = F.cross_entropy(sim_matrix, pos_indices)
    return loss

# テスト
torch.manual_seed(42)
model = SimCLR(input_dim=20, hidden_dim=64, proj_dim=32)
z1 = torch.randn(8, 32)
z2 = torch.randn(8, 32)
loss = info_nce_loss(z1, z2, temperature=0.5)
print(f"InfoNCE損失（ランダム初期化）: {loss.item():.4f}")
print(f"理論上のランダム時の損失 (log(2N-1)): {np.log(2*8 - 1):.4f}")

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.font_manager as _fm
for _c in ['Hiragino Sans','Yu Gothic','Noto Sans CJK JP','IPAexGothic','Meiryo']:
    if any(_c==_f.name for _f in _fm.fontManager.ttflist):
        plt.rcParams['font.family']=_c; break
plt.rcParams['axes.unicode_minus']=False

# 前のセクションで定義したクラスを再利用
class SimCLR(nn.Module):
    def __init__(self, input_dim, hidden_dim=64, proj_dim=32):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim), nn.BatchNorm1d(hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.BatchNorm1d(hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )
        self.projector = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, proj_dim)
        )

    def forward(self, x):
        h = self.encoder(x)
        z = self.projector(h)
        return h, z

    def get_representation(self, x):
        return self.encoder(x)

def info_nce_loss(z1, z2, temperature=0.5):
    batch_size = z1.size(0)
    z1 = F.normalize(z1, dim=1)
    z2 = F.normalize(z2, dim=1)
    z = torch.cat([z1, z2], dim=0)
    sim = torch.matmul(z, z.t()) / temperature
    mask = torch.eye(2 * batch_size, device=z.device).bool()
    sim.masked_fill_(mask, float('-inf'))
    pos_idx = torch.cat([torch.arange(batch_size, 2*batch_size),
                         torch.arange(0, batch_size)]).to(z.device)
    return F.cross_entropy(sim, pos_idx)

class SimpleAugmentation:
    def __init__(self, noise_scale=0.1, mask_prob=0.1):
        self.noise_scale = noise_scale
        self.mask_prob = mask_prob

    def __call__(self, x):
        noise = torch.randn_like(x) * self.noise_scale
        x_aug = x + noise
        mask = (torch.rand_like(x) > self.mask_prob).float()
        return x_aug * mask

# データ生成
np.random.seed(42)
num_classes = 5
dim = 20
data_list, labels_list = [], []
for c in range(num_classes):
    center = np.random.randn(dim) * 1.5
    samples = center + np.random.randn(200, dim) * 2.0
    data_list.append(samples)
    labels_list.append(np.full(200, c))
data = torch.FloatTensor(np.concatenate(data_list))
labels = torch.LongTensor(np.concatenate(labels_list))
perm = torch.randperm(len(data))
data, labels = data[perm], labels[perm]

# 訓練/テスト分割
split = int(len(data) * 0.8)
train_data, test_data = data[:split], data[split:]
train_labels, test_labels = labels[:split], labels[split:]

# SimCLRの訓練
torch.manual_seed(42)
model = SimCLR(input_dim=20, hidden_dim=64, proj_dim=32)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)
augmentation = SimpleAugmentation(noise_scale=0.3, mask_prob=0.15)

losses = []
epochs = 200
batch_size = 128

for epoch in range(epochs):
    model.train()
    # ランダムにバッチを取得
    idx = torch.randperm(len(train_data))[:batch_size]
    x_batch = train_data[idx]

    # 2つの拡張ビューを生成
    view1 = augmentation(x_batch)
    view2 = augmentation(x_batch)

    # 順伝播
    _, z1 = model(view1)
    _, z2 = model(view2)

    # InfoNCE損失
    loss = info_nce_loss(z1, z2, temperature=0.5)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    losses.append(loss.item())
    if (epoch + 1) % 50 == 0:
        print(f"Epoch {epoch+1}/{epochs}, Loss: {loss.item():.4f}")

# 訓練損失のプロット
fig, ax = plt.subplots(1, 1, figsize=(8, 5))
ax.plot(losses, color='#2ecc71', alpha=0.8)
ax.set_xlabel('エポック', fontsize=12)
ax.set_ylabel('InfoNCE損失', fontsize=12)
ax.set_title('SimCLRの学習損失', fontsize=14)
ax.axhline(y=np.log(2 * batch_size - 1), color='red', linestyle='--',
           alpha=0.5, label=f'ランダム基準 (log(2N-1)={np.log(2*batch_size-1):.2f})')
ax.legend(fontsize=10)
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('simclr_training_loss.png', dpi=150, bbox_inches='tight')
plt.show()

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.font_manager as _fm
for _c in ['Hiragino Sans','Yu Gothic','Noto Sans CJK JP','IPAexGothic','Meiryo']:
    if any(_c==_f.name for _f in _fm.fontManager.ttflist):
        plt.rcParams['font.family']=_c; break
plt.rcParams['axes.unicode_minus']=False
from sklearn.decomposition import PCA

# 表現の抽出
model.eval()
with torch.no_grad():
    train_repr = model.get_representation(train_data)
    test_repr = model.get_representation(test_data)

# 1. 線形プローブ: 凍結された表現の上に線形分類器を訓練
class LinearProbe(nn.Module):
    def __init__(self, input_dim, num_classes):
        super().__init__()
        self.linear = nn.Linear(input_dim, num_classes)

    def forward(self, x):
        return self.linear(x)

# 対照学習の表現で線形プローブ
probe_cl = LinearProbe(64, num_classes)
opt_cl = torch.optim.Adam(probe_cl.parameters(), lr=0.01)

# 比較: ランダム初期化のエンコーダの表現で線形プローブ
random_model = SimCLR(input_dim=20, hidden_dim=64, proj_dim=32)
with torch.no_grad():
    train_repr_random = random_model.get_representation(train_data)
    test_repr_random = random_model.get_representation(test_data)

probe_random = LinearProbe(64, num_classes)
opt_random = torch.optim.Adam(probe_random.parameters(), lr=0.01)

# 比較: 教師あり学習（直接学習）
supervised_model = nn.Sequential(
    nn.Linear(20, 64), nn.BatchNorm1d(64), nn.ReLU(),
    nn.Linear(64, 64), nn.BatchNorm1d(64), nn.ReLU(),
    nn.Linear(64, num_classes)
)
opt_sup = torch.optim.Adam(supervised_model.parameters(), lr=0.01)

results = {
    'SimCLR＋線形': [],
    'ランダム＋線形': [],
    '教師あり': []
}

for epoch in range(100):
    # SimCLR + Linear Probe
    probe_cl.train()
    logits = probe_cl(train_repr.detach())
    loss = F.cross_entropy(logits, train_labels)
    opt_cl.zero_grad(); loss.backward(); opt_cl.step()

    # Random + Linear Probe
    probe_random.train()
    logits_r = probe_random(train_repr_random.detach())
    loss_r = F.cross_entropy(logits_r, train_labels)
    opt_random.zero_grad(); loss_r.backward(); opt_random.step()

    # Supervised
    supervised_model.train()
    logits_s = supervised_model(train_data)
    loss_s = F.cross_entropy(logits_s, train_labels)
    opt_sup.zero_grad(); loss_s.backward(); opt_sup.step()

    if (epoch + 1) % 5 == 0:
        probe_cl.eval(); probe_random.eval(); supervised_model.eval()
        with torch.no_grad():
            acc_cl = (probe_cl(test_repr).argmax(1) == test_labels).float().mean().item()
            acc_r = (probe_random(test_repr_random).argmax(1) == test_labels).float().mean().item()
            acc_s = (supervised_model(test_data).argmax(1) == test_labels).float().mean().item()

        results['SimCLR＋線形'].append((epoch+1, acc_cl))
        results['ランダム＋線形'].append((epoch+1, acc_r))
        results['教師あり'].append((epoch+1, acc_s))

# 2. 表現のPCA可視化
fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# SimCLRの表現
pca = PCA(n_components=2)
repr_2d = pca.fit_transform(test_repr.detach().numpy())
scatter = axes[0].scatter(repr_2d[:, 0], repr_2d[:, 1], c=test_labels.numpy(),
                          cmap='tab10', s=20, alpha=0.7)
axes[0].set_title('SimCLRの表現（PCA）', fontsize=14)
axes[0].set_xlabel('第1主成分', fontsize=12)
axes[0].set_ylabel('第2主成分', fontsize=12)

# ランダムエンコーダの表現
repr_random_2d = pca.fit_transform(test_repr_random.detach().numpy())
axes[1].scatter(repr_random_2d[:, 0], repr_random_2d[:, 1], c=test_labels.numpy(),
                cmap='tab10', s=20, alpha=0.7)
axes[1].set_title('ランダムエンコーダ（PCA）', fontsize=14)
axes[1].set_xlabel('第1主成分', fontsize=12)
axes[1].set_ylabel('第2主成分', fontsize=12)

# テスト精度の比較
colors = {'SimCLR＋線形': '#2ecc71', 'ランダム＋線形': '#e74c3c', '教師あり': '#3498db'}
for name, res in results.items():
    epochs_list, accs = zip(*res)
    axes[2].plot(epochs_list, accs, 'o-', label=name, color=colors[name], markersize=3)
axes[2].set_xlabel('線形プローブのエポック', fontsize=12)
axes[2].set_ylabel('テスト精度', fontsize=12)
axes[2].set_title('下流分類の精度', fontsize=14)
axes[2].legend(fontsize=10)
axes[2].set_ylim(0, 1.05)
axes[2].grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('simclr_evaluation.png', dpi=150, bbox_inches='tight')
plt.show()

# 最終精度
for name, res in results.items():
    print(f"{name:>20s}: {res[-1][1]*100:.1f}%")

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.font_manager as _fm
for _c in ['Hiragino Sans','Yu Gothic','Noto Sans CJK JP','IPAexGothic','Meiryo']:
    if any(_c==_f.name for _f in _fm.fontManager.ttflist):
        plt.rcParams['font.family']=_c; break
plt.rcParams['axes.unicode_minus']=False

# 再度モデルとデータの定義（前セクションと同一）
class SimCLRTemp(nn.Module):
    def __init__(self, input_dim, hidden_dim=64, proj_dim=32):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim), nn.BatchNorm1d(hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.BatchNorm1d(hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )
        self.projector = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, proj_dim)
        )

    def forward(self, x):
        h = self.encoder(x)
        z = self.projector(h)
        return h, z

    def get_representation(self, x):
        return self.encoder(x)

def info_nce_loss(z1, z2, temperature=0.5):
    bs = z1.size(0)
    z1 = F.normalize(z1, dim=1)
    z2 = F.normalize(z2, dim=1)
    z = torch.cat([z1, z2], dim=0)
    sim = torch.matmul(z, z.t()) / temperature
    mask = torch.eye(2*bs, device=z.device).bool()
    sim.masked_fill_(mask, float('-inf'))
    pos_idx = torch.cat([torch.arange(bs, 2*bs), torch.arange(0, bs)]).to(z.device)
    return F.cross_entropy(sim, pos_idx)

class SimpleAug:
    def __init__(self, noise_scale=0.3, mask_prob=0.15):
        self.noise_scale = noise_scale
        self.mask_prob = mask_prob
    def __call__(self, x):
        x = x + torch.randn_like(x) * self.noise_scale
        return x * (torch.rand_like(x) > self.mask_prob).float()

# データ生成
np.random.seed(42)
data_list, labels_list = [], []
for c in range(5):
    center = np.random.randn(20) * 1.5
    data_list.append(center + np.random.randn(200, 20) * 2.0)
    labels_list.append(np.full(200, c))
data = torch.FloatTensor(np.concatenate(data_list))
labels = torch.LongTensor(np.concatenate(labels_list))
perm = torch.randperm(len(data))
data, labels = data[perm], labels[perm]
split = int(len(data) * 0.8)
train_data, test_data = data[:split], data[split:]
train_labels, test_labels = labels[:split], labels[split:]

# 異なる温度での訓練
temperatures = [0.05, 0.1, 0.5, 1.0, 2.0]
aug = SimpleAug()
temp_results = {}

for tau in temperatures:
    torch.manual_seed(42)
    model = SimCLRTemp(20, 64, 32)
    opt = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)
    losses = []

    for epoch in range(200):
        model.train()
        idx = torch.randperm(len(train_data))[:128]
        v1, v2 = aug(train_data[idx]), aug(train_data[idx])
        _, z1 = model(v1)
        _, z2 = model(v2)
        loss = info_nce_loss(z1, z2, temperature=tau)
        opt.zero_grad(); loss.backward(); opt.step()
        losses.append(loss.item())

    # 線形プローブで評価
    model.eval()
    with torch.no_grad():
        train_repr = model.get_representation(train_data)
        test_repr = model.get_representation(test_data)

    probe = nn.Linear(64, 5)
    opt_p = torch.optim.Adam(probe.parameters(), lr=0.01)
    for ep in range(100):
        probe.train()
        logits = probe(train_repr.detach())
        l = F.cross_entropy(logits, train_labels)
        opt_p.zero_grad(); l.backward(); opt_p.step()

    probe.eval()
    with torch.no_grad():
        acc = (probe(test_repr).argmax(1) == test_labels).float().mean().item()

    temp_results[tau] = {'losses': losses, 'accuracy': acc}
    print(f"tau={tau:.2f}: テスト精度={acc*100:.1f}%")

# 可視化
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# 損失の推移
colors_temp = plt.cm.viridis(np.linspace(0, 1, len(temperatures)))
for (tau, res), color in zip(temp_results.items(), colors_temp):
    axes[0].plot(res['losses'], label=f'tau={tau}', color=color, alpha=0.8)
axes[0].set_xlabel('エポック', fontsize=12)
axes[0].set_ylabel('InfoNCE損失', fontsize=12)
axes[0].set_title('温度別の学習損失', fontsize=14)
axes[0].legend(fontsize=10)
axes[0].grid(True, alpha=0.3)

# 精度の比較
taus = list(temp_results.keys())
accs = [temp_results[t]['accuracy'] for t in taus]
bars = axes[1].bar([str(t) for t in taus], accs, color='#2ecc71', edgecolor='black', linewidth=0.8)
axes[1].set_xlabel('温度 tau', fontsize=12)
axes[1].set_ylabel('テスト精度（線形プローブ）', fontsize=12)
axes[1].set_title('温度が下流精度に与える影響', fontsize=14)
axes[1].set_ylim(0, 1.05)
for bar, acc in zip(bars, accs):
    axes[1].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,
                f'{acc*100:.1f}%', ha='center', fontsize=11, fontweight='bold')
axes[1].grid(True, alpha=0.3, axis='y')

plt.tight_layout()
plt.savefig('temperature_effect.png', dpi=150, bbox_inches='tight')
plt.show()

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.font_manager as _fm
for _c in ['Hiragino Sans','Yu Gothic','Noto Sans CJK JP','IPAexGothic','Meiryo']:
    if any(_c==_f.name for _f in _fm.fontManager.ttflist):
        plt.rcParams['font.family']=_c; break
plt.rcParams['axes.unicode_minus']=False

# 学習済みモデルと未学習モデルの類似度行列を比較
# ここでは小さなバッチで可視化

class SimCLRViz(nn.Module):
    def __init__(self, input_dim, hidden_dim=64, proj_dim=32):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim), nn.BatchNorm1d(hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.BatchNorm1d(hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )
        self.projector = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, proj_dim)
        )
    def forward(self, x):
        return self.projector(self.encoder(x))

def info_nce_loss_fn(z1, z2, temperature=0.5):
    bs = z1.size(0)
    z1 = F.normalize(z1, dim=1)
    z2 = F.normalize(z2, dim=1)
    z = torch.cat([z1, z2], dim=0)
    sim = torch.matmul(z, z.t()) / temperature
    mask = torch.eye(2*bs, device=z.device).bool()
    sim.masked_fill_(mask, float('-inf'))
    pos_idx = torch.cat([torch.arange(bs, 2*bs), torch.arange(0, bs)]).to(z.device)
    return F.cross_entropy(sim, pos_idx)

class SimpleAugViz:
    def __init__(self, ns=0.3, mp=0.15):
        self.ns, self.mp = ns, mp
    def __call__(self, x):
        return (x + torch.randn_like(x)*self.ns) * (torch.rand_like(x)>self.mp).float()

# データ生成
np.random.seed(42)
data_list, labels_list = [], []
for c in range(5):
    center = np.random.randn(20) * 1.5
    data_list.append(center + np.random.randn(200, 20) * 2.0)
    labels_list.append(np.full(200, c))
data = torch.FloatTensor(np.concatenate(data_list))
labels = torch.LongTensor(np.concatenate(labels_list))
perm = torch.randperm(len(data))
data, labels = data[perm], labels[perm]

# 各クラスから2サンプルずつ選択（計10サンプル）
selected_indices = []
for c in range(5):
    class_indices = (labels == c).nonzero(as_tuple=True)[0][:2]
    selected_indices.extend(class_indices.tolist())
selected_data = data[selected_indices]
selected_labels = labels[selected_indices]

# 学習
torch.manual_seed(42)
model_trained = SimCLRViz(20, 64, 32)
aug = SimpleAugViz()
opt = torch.optim.Adam(model_trained.parameters(), lr=0.001)

for epoch in range(300):
    model_trained.train()
    idx = torch.randperm(len(data))[:128]
    v1, v2 = aug(data[idx]), aug(data[idx])
    loss = info_nce_loss_fn(model_trained(v1), model_trained(v2))
    opt.zero_grad(); loss.backward(); opt.step()

# 未学習モデル
torch.manual_seed(0)
model_random = SimCLRViz(20, 64, 32)

# 類似度行列の計算と可視化
fig, axes = plt.subplots(1, 2, figsize=(14, 6))

for ax, model, title in [(axes[0], model_random, '学習前（ランダム）'),
                          (axes[1], model_trained, '対照学習の後')]:
    model.eval()
    with torch.no_grad():
        z = model(selected_data)
        z = F.normalize(z, dim=1)
        sim = torch.matmul(z, z.t()).numpy()

    im = ax.imshow(sim, cmap='RdBu_r', vmin=-1, vmax=1)
    ax.set_title(title, fontsize=14)

    # ラベル表示
    tick_labels = [f'C{selected_labels[i].item()}' for i in range(len(selected_labels))]
    ax.set_xticks(range(len(tick_labels)))
    ax.set_xticklabels(tick_labels, fontsize=10)
    ax.set_yticks(range(len(tick_labels)))
    ax.set_yticklabels(tick_labels, fontsize=10)

    # 数値を表示
    for i in range(len(sim)):
        for j in range(len(sim)):
            ax.text(j, i, f'{sim[i,j]:.2f}', ha='center', va='center',
                    fontsize=8, color='white' if abs(sim[i,j]) > 0.5 else 'black')

    plt.colorbar(im, ax=ax, fraction=0.046)

plt.suptitle('コサイン類似度行列：各クラス2サンプル', fontsize=15, y=1.02)
plt.tight_layout()
plt.savefig('similarity_matrix_comparison.png', dpi=150, bbox_inches='tight')
plt.show()