<jupyter_start><jupyter_text><jupyter_code>import torch
import torch.nn as nn 
import numpy as np
import torch.nn.functional as F<jupyter_output><empty_output><jupyter_text>affinity lossの定義<jupyter_code>class AffinityLoss(nn.Module):
  def __init__(self, n_classes, dim, sigma, xi):
    super(AffinityLoss, self).__init__()
    self.n_classes = n_classes
    self.dim = dim 
    self.sigma = sigma
    self.xi = xi
    self.w = nn.parameter.Parameter(torch.empty(self.n_classes, self.dim))
    nn.init.xavier_normal_(self.w)
  
  def forward(self, f, label):
    # calcurate gaussian similarity
    # d(f_i, w_j) = sim[i + n_classes * j]
    f_expand = torch.cat([f] * self.n_classes, dim=0)
    w_expand = torch.cat([self.w] * f.size(0), dim=0)
    l2_norm = torch.sum((f_expand - w_expand)**2, 1)
    d = torch.exp(-l2_norm / self.sigma)

    label_one_hot = F.one_hot(label, num_classes=self.n_classes)
    label_flatten = torch.flatten(label_one_hot)
    L_hat = self.xi + (1 - label_flatten) * d - label_flatten * d 
    L_pos = torch.where(L_hat > 0, L_hat, torch.zeros_like(L_hat))
    L_mm = torch.sum(L_pos, dim=0)

    # diversity legularizar
    W = torch.matmul(self.w, self.w.transpose(1, 0))
    W_upper_triangle = torch.triu(W, 1)
    mu = 2 / (self.n_classes**2 - self.n_classes) * torch.sum(W_upper_triangle)
    X = W - mu
    X_upper_triangle = torch.triu(X, 1)
    Rw = torch.sum(X_upper_triangle**2)

    return L_mm + Rw<jupyter_output><empty_output><jupyter_text>モデルの定義<jupyter_code>from torchvision import models

class FeatureExtracter(nn.Module):
  def __init__(self):
    super(FeatureExtracter, self).__init__()
    self.base_model = models.mobilenet_v2(pretrained=True).features

    self.gap = nn.AdaptiveAvgPool2d((1, 1))
    self.fc = nn.Linear(1280, 512)
    self.act = nn.ReLU()

  def forward(self, input):
    x = self.base_model(input)
    x = self.gap(x)
    x = x.view(-1, self.num_flat_features(x))
    x = self.fc(x)
    output = self.act(x)
    return output 
  
  def num_flat_features(self, x):
    size = x.size()[1:]  # all dimensions except the batch dimension
    num_features = 1
    for s in size:
      num_features *= s
    return num_features


model = FeatureExtracter()<jupyter_output><empty_output><jupyter_text>学習データのダウンロード<jupyter_code>import torchvision
from torchvision import transforms

batch_size=64

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])

transform = transforms.Compose([transforms.ToTensor(), normalize])

trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=0)

testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=0)<jupyter_output>Files already downloaded and verified
Files already downloaded and verified<jupyter_text>損失関数、オプティマイザの定義<jupyter_code>from torch import optim

classes = trainset.classes
n_classes = len(classes)

criterion = AffinityLoss(n_classes, 512, 10, 0.5)

optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)<jupyter_output><empty_output><jupyter_text>学習<jupyter_code>device = 'cuda'

model.to(device)
criterion.to(device)

for epoch in range(20):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        inputs = inputs.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        # backpropagation
        loss.backward()
        optimizer.step()
        train_loss = loss.item()
        running_loss += loss.item()
        if i % batch_size == batch_size - 1:
            print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / batch_size))
            running_loss = 0.0
print('Finished Training')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12641138929798723, 'fraction_numerical': 0.021600392734413353, 'mean_word_length': 3.811097992916175, 'pattern_counts': {'":': 0, '<': 24, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 5, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '107965', 'license': 'Apache-2.0', 'n_tokens_mistral': 1575, 'n_tokens_neox': 1497, 'n_words': 281, 'path': 'Untitled0.ipynb', 'repo_name': 'yuu399/AffinityLoss', 'size': '4074'}