iOSで動作する異常検知モデルを作った

from torchvision.models import MobileNetV2

class MobileNetFeatures(nn.Module):
    def __init__(self):
        super(MobileNetFeatures, self).__init__()
        self.head = MobileNetV2().features
        self.pool = nn.AvgPool2d(4, 4)
        self.flat = nn.Flatten()
        self.fc = nn.Linear(1280, 512)

    def forward(self, x):
        x = self.head(x)
        x = self.pool(x)
        x = self.flat(x)
        x = self.fc(x)
        return x

    if save_interval > 0 and epoch_id % save_interval == 0:
        model.eval()
　　　　　# 正常画像と異常画像のコサイン類似度をそれぞれ測定する.
        positive_dist = []
        negative_dist = []
        for batch in valid_loader:
            images = batch[0].to(device)
            labels = batch[1].numpy().tolist()
            labels = [bool(i) for i in labels]
            with torch.no_grad():
                embeddings = model(images).cpu().numpy()

            positive_embeddings = embeddings[labels]
            negative_embeddings = embeddings[[not i for i in labels]]

            mean_embedding = np.mean(positive_embeddings, axis=0)
            for pe in positive_embeddings:
                cos_sim = np.dot(mean_embedding, pe) / (np.linalg.norm(mean_embedding, ord=2) * np.linalg.norm(pe, ord=2))
                positive_dist.append(cos_sim)
            for ne in negative_embeddings:
                cos_sim = np.dot(mean_embedding, ne) / (np.linalg.norm(mean_embedding, ord=2) * np.linalg.norm(ne, ord=2))
                negative_dist.append(cos_sim)
        mean_positive_dist = sum(positive_dist) / len(positive_dist)
        mean_negative_dist = sum(negative_dist) / len(negative_dist)

        print(f"epoch{epoch_id}: {mean_positive_dist} {mean_negative_dist}")
        model.train()

        # embeddingを保存する
        features_save_path = f"../saved_features/embedding.txt"
        np.savetxt(features_save_path, mean_embedding, delimiter=",")