Transformer PyTorch 多分类器多指标小 demo

import torch
import numpy as np
from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score, roc_auc_score

# 构建Transformer模型
class MyTransformer(torch.nn.Module):
    def __init__(self):
        super(MyTransformer, self).__init__()
        self.embedding = torch.nn.Linear(5, 16)
        self.transformer1 = torch.nn.TransformerEncoderLayer(16, nhead=4)
        self.transformer2 = torch.nn.TransformerEncoderLayer(16, nhead=4)
        self.transformer3 = torch.nn.TransformerEncoderLayer(16, nhead=4)
        self.fc = torch.nn.Linear(16, 3)
        self.softmax = torch.nn.Softmax(dim=1)

    def forward(self, x):
        x = self.embedding(x)
        x = x.permute(1, 0, 2)  # 调整维度顺序，使得时间步维度在第一维
        x = self.transformer1(x)
        x = self.transformer2(x)
        x = self.transformer3(x)
        x = x.permute(1, 0, 2)  # 调整维度顺序，使得时间步维度在第二维
        x = self.fc(x[:, -1, :])  # 取最后一个时间步的输出，作为整个序列的输出
        return x

# 构造训练数据
X_train = torch.randn(100, 10, 5)  # 100个样本，每个样本10个时间步，每个时间步5个特征
y_train = torch.randint(0, 3, (100, ))  # 3个类别

# 构建Transformer模型
model = MyTransformer()

# 训练模型
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
for epoch in range(100):
    optimizer.zero_grad()
    outputs = model(X_train)
    loss = criterion(outputs, y_train)
    loss.backward()
    optimizer.step()

# 构造测试数据
X_test = torch.randn(10, 10, 5)  # 10个样本，每个样本10个时间步，每个时间步5个特征
y_test = torch.randint(0, 3, (10, ))  # 3个类别

# 预测
model.eval()
with torch.no_grad():
    y_pred = model(X_test)
    y_scores = y_pred.numpy()

# 计算p, r, f1, acc
y_pred = np.argmax(y_pred.numpy(), axis=1)
p = precision_score(y_test, y_pred, average='macro')
r = recall_score(y_test, y_pred, average='macro')
f1 = f1_score(y_test, y_pred, average='macro')
acc = accuracy_score(y_test, y_pred)

# 计算auc
y_prob = torch.softmax(torch.tensor(y_scores), dim=1).numpy()
auc = roc_auc_score(y_test, y_prob, multi_class='ovr')

# 输出指标
print('Precision: {:.4f}'.format(p))
print('Recall: {:.4f}'.format(r))
print('F1: {:.4f}'.format(f1))
print('Accuracy: {:.4f}'.format(acc))
print('AUC: {:.4f}'.format(auc))