ArcFaceLoss和CenterLoss的代碼實現(xiàn)

手動實現(xiàn)ArcFaceLoss和CenterLoss，并用來訓(xùn)練MNIST數(shù)據(jù)。

導(dǎo)入相關(guān)庫
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms as T
from torch.utils.data import DataLoader
import itertools
import matplotlib.pyplot as plt
# 查看時間和進度
from tqdm import tqdm
import time
實現(xiàn)ArcFaceNet和CenterLossNet
class ArcFaceNet(nn.Module):
    def __init__(self, cls_num=10, feature_dim=2):
        super(ArcFaceNet, self).__init__()
        self.w = nn.Parameter(torch.randn(feature_dim, cls_num))

    def forward(self, features, m=1, s=10):
        # 特征與權(quán)重 歸一化
        _features = nn.functional.normalize(features, dim=1)
        _w = nn.functional.normalize(self.w, dim=0)

        # 特征向量與參數(shù)向量的夾角theta，分子numerator，分母denominator
        theta = torch.acos(torch.matmul(_features, _w) / 10)  # /10防止下溢
        numerator = torch.exp(s * torch.cos(theta + m))
        denominator = torch.sum(torch.exp(s * torch.cos(theta)), dim=1, keepdim=True) - torch.exp(
            s * torch.cos(theta)) + numerator
        return torch.log(torch.div(numerator, denominator))
    
class CenterLossNet(nn.Module):
    def __init__(self, cls_num=10, feature_dim=2):
        super(CenterLossNet, self).__init__()
        self.centers = nn.Parameter(torch.randn(cls_num, feature_dim))

    def forward(self, features, labels, reduction='mean'):
        # 特征向量歸一化
        _features = nn.functional.normalize(features)

        centers_batch = self.centers.index_select(dim=0, index=labels.long())
        # 根據(jù)論文《A Discriminative Feature Learning Approach for Deep Face Recognition》修改如下
        if reduction == 'sum':  # 返回loss的和
            return torch.sum(torch.pow(_features - centers_batch, 2)) / 2
        elif reduction == 'mean':  # 返回loss和的平均值，默認為mean方式
            return torch.sum(torch.pow(_features - centers_batch, 2)) / 2 / len(features)
        else:
            raise ValueError("ValueError: {0} is not a valid value for reduction".format(reduction))
定義LeNet模型
class LeNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 64, 3, padding=1),
            nn.PReLU(),
            nn.BatchNorm2d(64),
            nn.Conv2d(64, 32, 3, stride=2, padding=1),
            nn.PReLU(),
            nn.BatchNorm2d(32),
            nn.modules.Flatten()
        )
        self.linear = nn.Sequential(
            nn.Linear(32 * 14 * 14, 512),
            nn.PReLU(),
            nn.BatchNorm1d(512),
            nn.Linear(512, 256),
            nn.PReLU(),
            nn.BatchNorm1d(256),
            nn.Linear(256, 64),
            nn.PReLU(),
            nn.BatchNorm1d(64),
            nn.Linear(64, 32)
            # nn.Linear(64, 2, bias=False)  # features設(shè)置為二維，可以進行可視化
        )
        self.out_layer = nn.Sequential(
            nn.Linear(32, 10),
            # nn.Linear(2, 10),  # features設(shè)置為二維，可以進行可視化
            nn.LogSoftmax(dim=1)  # LogSoftmax與net=nn.NLLLoss()結(jié)合使用,求交叉熵損失
        )

    def forward(self, x):
        features = self.linear(self.conv(x))
        out = self.out_layer(features)  # 用于計算CrossEntropyLoss
        return features, out
模型訓(xùn)練

兩種損失計算方式：

  • CrossEntropyLoss+CenterLoss
  • ArcFaceLoss+CenterLoss

超參數(shù)都是初始隨便設(shè)定的，跑了一遍，精度可達到99.29。你可以調(diào)調(diào)超參數(shù)，精度可以更高。訓(xùn)練代碼如下：

# 特征向量可視化
def visualize(features, labels, loss, epoch):
    # 定義10種顏色
    colors = ['#ff0000', '#ffff00', '#00ff00', '#00ffff', '#0000ff', '#ff00ff', '#990000', '#999900', '#009900',
              '#009999']

    plt.clf()  # 清空畫板
    # 畫出所有的點，不同的label對應(yīng)不同的顏色
    for i in range(10):
        plt.plot(features[labels == i, 0], features[labels == i, 1], ".", c=colors[i], label=i)
    plt.legend(loc="upper right")  # 圖例
    plt.title(f"ce+cl: epoch={epoch}, loss={loss}")  # 標題
    plt.savefig("ce+cl/image/epoch%d.jpg" % epoch)  # 保存圖片
    plt.draw()  # 展示圖片
    plt.pause(0.001)

# 1.加載數(shù)據(jù)集
transform_op = T.Compose([  # 數(shù)據(jù)預(yù)處理
    T.ToTensor(),
    T.Normalize([0.4914], [0.2023])
])
train_dataset = datasets.MNIST("../code/data", train=True, transform=transform_op, download=False)
val_dataset = datasets.MNIST("../code/data", train=False, transform=transform_op, download=False)
train_dataloader = DataLoader(train_dataset, batch_size=128, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=128, shuffle=False)

# 訓(xùn)練設(shè)備: GPU or CPU
device= torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 超參
lr = 1e-3
epochs = 20
lr_schedule = {
    5: 1e-3,
    10: 1e-4,
    15: 1e-5
}
alpha = 0.95  # centerloss與arcfaceloss的權(quán)重比例

1.CrossEntropyLoss+CenterLoss

# 2.創(chuàng)建模型
cls_num, feature_dim = 10, 32  # 10分類
# cls_num, feature_dim = 10, 2  # features設(shè)置為二維，可以進行可視化
net = LeNet().to(device)
centerloss_net = CenterLossNet(cls_num, feature_dim).to(device)
# 3.定義損失
loss_func = nn.NLLLoss()
# 4.定義優(yōu)化器
optimizer = optim.Adam(itertools.chain(net.parameters(), centerloss_net.parameters()), lr)

# 5.模型訓(xùn)練
plt.ion()
for epoch in range(epochs):
    start = time.time()
    # 學(xué)習(xí)率策略
    if epoch in lr_schedule:
        lr = lr_schedule[epoch]
        for group in optimizer.param_groups:
            group["lr"] = lr
            
    # 1）訓(xùn)練集
    net.train()  # train mode
    features_loader, labels_loader = [], []  # 保存特征向量和標簽的列表，用于可視化操作
    train_loss = 0.
    for images, targets in tqdm(train_dataloader):
        images, targets = images.to(device), targets.to(device)
        
        # 方式1: CrossEntropyLoss+CenterLoss
        features, out = net(images)
        # 計算損失
        ce_loss = loss_func(out, targets)
        center_loss = centerloss_net(features, targets)
        loss = alpha * ce_loss + (1 - alpha) * center_loss
        
        optimizer.zero_grad()  # 清空梯度
        loss.backward()  # 反向傳播
        optimizer.step()  # 梯度更新
        
        # 統(tǒng)計訓(xùn)練損失
        train_loss += loss.cpu().detach().item()

        # 將特征和標簽加入到列表中
        features_loader.append(features)
        labels_loader.append(targets)
    # 計算平均損失
    train_loss /= len(train_dataloader)

    # 2.測試集
    net.eval()  # evaluation mode
    val_loss, correct = 0., 0.
    with torch.no_grad():  # 作用域范圍內(nèi)不計算梯度，節(jié)省內(nèi)存
        for images, targets in tqdm(val_dataloader):
            images, targets = images.to(device), targets.to(device)
            
            # 方式1: CrossEntropyLoss+CenterLoss
            features, out = net(images)
            # 計算損失
            ce_loss = loss_func(out, targets)
            center_loss = centerloss_net(features, targets)
            loss = alpha * ce_loss + (1 - alpha) * center_loss

            # 統(tǒng)計驗證損失
            val_loss += loss.cpu().detach().item()
            # 統(tǒng)計正確的個數(shù)
            correct += sum(out.argmax(1) == targets)
        # 計算平均損失
        val_loss /= len(val_dataloader)
    # 計算準確率
    accuracy = correct.item() / len(val_dataset)
    
    # 打印損失和精度信息
    print(f"Epoch: {epoch}/{epochs}, Train_loss: {train_loss:.5f}, Val_loss: {val_loss:.5f}, Accuracy: {accuracy}")
    # 保存模型參數(shù)
    torch.save(net.state_dict(), f"ce+cl/checkpoint/net.pt")
    torch.save(centerloss_net.state_dict(), f"ce+cl/checkpoint/centerloss_net.pt")
    # 特征向量可視化
    features = torch.cat(features_loader, dim=0)
    labels = torch.cat(labels_loader, dim=0)
    visualize(features.cpu().detach().numpy(), labels.cpu().detach().numpy(), train_loss, epoch)
    # 查看時間和進度
    end = time.time()  # 本次輪詢結(jié)束時間
    print(f"第{epoch}次輪詢，共耗時{int(end - start)}秒")
    time.sleep(0.01)
plt.ioff()

2.ArcFaceLoss+CenterLoss

# 2.創(chuàng)建模型
cls_num, feature_dim = 10, 32  # 10分類
# cls_num, feature_dim = 10, 2  # features設(shè)置為二維，可以進行可視化
net = LeNet().to(device)
arcface_net = ArcFaceNet(cls_num, feature_dim).to(device)
centerloss_net = CenterLossNet(cls_num, feature_dim).to(device)
# 3.定義損失
loss_func = nn.NLLLoss()
# 4.定義優(yōu)化器
optimizer = optim.Adam(itertools.chain(net.parameters(), arcface_net.parameters(), centerloss_net.parameters()), lr)

# 5.模型訓(xùn)練
plt.ion()
for epoch in range(epochs):
    start = time.time()
    # 學(xué)習(xí)率策略
    if epoch in lr_schedule:
        lr = lr_schedule[epoch]
        for group in optimizer.param_groups:
            group["lr"] = lr
            
    # 1）訓(xùn)練集
    net.train()  # train mode
    features_loader, labels_loader = [], []  # 保存特征向量和標簽的列表，用于可視化操作
    train_loss = 0.
    for images, targets in tqdm(train_dataloader):
        images, targets = images.to(device), targets.to(device)

        # 方式2: ArcFaceLoss+CenterLoss
        features, _ = net(images)
        out = arcface_net(features)
        # 計算損失
        arcface_loss = loss_func(out, targets)  # arcfaceloss
        center_loss = centerloss_net(features, targets)  # centerloss
        loss = alpha * arcface_loss + (1 - alpha) * center_loss
        
        optimizer.zero_grad()  # 清空梯度
        loss.backward()  # 反向傳播
        optimizer.step()  # 梯度更新
        
        # 統(tǒng)計訓(xùn)練損失
        train_loss += loss.cpu().detach().item()

        # 將特征和標簽加入到列表中
        features_loader.append(features)
        labels_loader.append(targets)
    # 計算平均損失
    train_loss /= len(train_dataloader)

    # 2.測試集
    net.eval()  # evaluation mode
    val_loss, correct = 0., 0.
    with torch.no_grad():  # 作用域范圍內(nèi)不計算梯度，節(jié)省內(nèi)存
        for images, targets in tqdm(val_dataloader):
            images, targets = images.to(device), targets.to(device)

            # 方式2: ArcFaceLoss+CenterLoss
            features, _ = net(images)
            out = arcface_net(features)
            # 計算損失
            arcface_loss = loss_func(out, targets)  # arcfaceloss
            center_loss = centerloss_net(features, targets)  # centerloss
            loss = alpha * arcface_loss + (1 - alpha) * center_loss

            # 統(tǒng)計驗證損失
            val_loss += loss.cpu().detach().item()
            # 統(tǒng)計正確的個數(shù)
            correct += sum(out.argmax(1) == targets)
        # 計算平均損失
        val_loss /= len(val_dataloader)
    # 計算準確率
    accuracy = correct.item() / len(val_dataset)
    
    # 打印損失和精度信息
    print(alpha * arcface_loss, (1 - alpha) * center_loss, arcface_loss, center_loss)
    print(f"Epoch: {epoch}/{epochs}, Train_loss: {train_loss:.5f}, Val_loss: {val_loss:.5f}, Accuracy: {accuracy}")
    # 保存模型參數(shù)
    torch.save(net.state_dict(), f"arcface+cl/checkpoint/net.pt")
    torch.save(centerloss_net.state_dict(), f"arcface+cl/checkpoint/centerloss_net.pt")
    torch.save(arcface_net.state_dict(), f"arcface+cl/checkpoint/arcface_net.pt")
    # 特征向量可視化
    features = torch.cat(features_loader, dim=0)
    labels = torch.cat(labels_loader, dim=0)
    visualize(features.cpu().detach().numpy(), labels.cpu().detach().numpy(), epoch, train_loss, val_loss, accuracy)
    # 查看時間和進度
    end = time.time()  # 本次輪詢結(jié)束時間
    print(f"第{epoch}次輪詢，共耗時{int(end - start)}秒")
    time.sleep(0.01)
plt.ioff()
最后編輯于 ：
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成，瀏覽時請結(jié)合常識與多方信息審慎甄別。
平臺聲明：文章內(nèi)容（如有圖片或視頻亦包括在內(nèi)）由作者上傳并發(fā)布，文章內(nèi)容僅代表作者本人觀點，簡書系信息發(fā)布平臺，僅提供信息存儲服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容