初始提交
This commit is contained in:
268
classification_model/Classification/CNN_Transfomer.py
Normal file
268
classification_model/Classification/CNN_Transfomer.py
Normal file
@ -0,0 +1,268 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
from torch.cuda.amp import GradScaler, autocast
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
from sklearn.preprocessing import StandardScaler
|
||||
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix
|
||||
from sklearn.utils.class_weight import compute_class_weight
|
||||
from sklearn.model_selection import train_test_split
|
||||
import numpy as np
|
||||
import os
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
|
||||
|
||||
# 自定义数据集
|
||||
class MyDataset(Dataset):
|
||||
def __init__(self, specs, labels, augment=False):
|
||||
self.specs = specs
|
||||
self.labels = labels
|
||||
self.augment = augment
|
||||
|
||||
def __getitem__(self, index):
|
||||
spec, target = self.specs[index], self.labels[index]
|
||||
if self.augment:
|
||||
noise = 0.01 * torch.randn_like(spec)
|
||||
spec = spec + noise
|
||||
return spec, target
|
||||
|
||||
def __len__(self):
|
||||
return len(self.specs)
|
||||
|
||||
# 数据标准化
|
||||
def ZspProcess(X_train, X_test, y_train, y_test, need=True):
|
||||
if need:
|
||||
scaler = StandardScaler()
|
||||
X_train = scaler.fit_transform(X_train)
|
||||
X_test = scaler.transform(X_test)
|
||||
|
||||
X_train = torch.tensor(X_train[:, np.newaxis, :], dtype=torch.float32)
|
||||
X_test = torch.tensor(X_test[:, np.newaxis, :], dtype=torch.float32)
|
||||
y_train = torch.tensor(y_train, dtype=torch.long)
|
||||
y_test = torch.tensor(y_test, dtype=torch.long)
|
||||
|
||||
data_train = MyDataset(X_train, y_train, augment=True)
|
||||
data_test = MyDataset(X_test, y_test, augment=False)
|
||||
return data_train, data_test
|
||||
|
||||
# Focal Loss
|
||||
class FocalLoss(nn.Module):
|
||||
def __init__(self, alpha=1, gamma=2, reduction='mean'):
|
||||
super(FocalLoss, self).__init__()
|
||||
self.alpha = alpha
|
||||
self.gamma = gamma
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, inputs, targets):
|
||||
probs = torch.softmax(inputs, dim=1)
|
||||
target_probs = probs[range(len(targets)), targets]
|
||||
focal_weight = self.alpha * (1 - target_probs) ** self.gamma
|
||||
log_prob = -torch.log(target_probs)
|
||||
loss = focal_weight * log_prob
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
else:
|
||||
return loss
|
||||
|
||||
# Transformer模块
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1):
|
||||
super(TransformerBlock, self).__init__()
|
||||
self.attention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout)
|
||||
self.feed_forward = nn.Sequential(
|
||||
nn.Linear(embed_dim, ff_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(ff_dim, embed_dim)
|
||||
)
|
||||
self.layernorm1 = nn.LayerNorm(embed_dim)
|
||||
self.layernorm2 = nn.LayerNorm(embed_dim)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, x):
|
||||
attn_output, _ = self.attention(x, x, x)
|
||||
x = self.layernorm1(x + self.dropout(attn_output))
|
||||
ff_output = self.feed_forward(x)
|
||||
x = self.layernorm2(x + self.dropout(ff_output))
|
||||
return x
|
||||
|
||||
# 改进后的CNN+Transformer模型
|
||||
class CNNWithTransformer(nn.Module):
|
||||
def __init__(self, nls, embed_dim=96, num_heads=2, ff_dim=192, dropout=0.1):
|
||||
super(CNNWithTransformer, self).__init__()
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.Conv1d(1, 64, kernel_size=5, stride=2, padding=2),
|
||||
nn.BatchNorm1d(64),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.2), # 添加Dropout
|
||||
nn.MaxPool1d(2, 2)
|
||||
)
|
||||
self.conv2 = nn.Sequential(
|
||||
nn.Conv1d(64, embed_dim, kernel_size=5, stride=2, padding=2),
|
||||
nn.BatchNorm1d(embed_dim),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.2), # 添加Dropout
|
||||
nn.MaxPool1d(2, 2)
|
||||
)
|
||||
self.transformer = TransformerBlock(embed_dim, num_heads, ff_dim, dropout)
|
||||
self.fc = nn.Sequential(
|
||||
nn.Linear(embed_dim, 128),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3), # 添加Dropout
|
||||
nn.Linear(128, nls)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.conv2(x)
|
||||
x = x.permute(2, 0, 1)
|
||||
x = self.transformer(x)
|
||||
x = x.mean(dim=0)
|
||||
x = self.fc(x)
|
||||
return x
|
||||
|
||||
# 训练函数（包含早停机制）
|
||||
def TransformerTrain(X_train, X_val, y_train, y_val, BATCH_SIZE, n_epochs, nls, model_path, patience=10):
|
||||
data_train, data_val = ZspProcess(X_train, X_val, y_train, y_val, need=True)
|
||||
train_loader = DataLoader(data_train, batch_size=BATCH_SIZE, shuffle=True)
|
||||
val_loader = DataLoader(data_val, batch_size=BATCH_SIZE, shuffle=False)
|
||||
|
||||
model = CNNWithTransformer(nls=nls).to(device)
|
||||
optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=0.001)
|
||||
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5)
|
||||
criterion = FocalLoss(alpha=1, gamma=2).to(device)
|
||||
scaler = GradScaler()
|
||||
|
||||
best_val_loss = float('inf')
|
||||
early_stop_counter = 0
|
||||
y_true_train, y_pred_train = [], []
|
||||
|
||||
for epoch in range(n_epochs):
|
||||
model.train()
|
||||
train_loss, train_acc = [], []
|
||||
|
||||
for inputs, labels in train_loader:
|
||||
inputs, labels = inputs.to(device), labels.to(device)
|
||||
optimizer.zero_grad()
|
||||
|
||||
with autocast():
|
||||
outputs = model(inputs)
|
||||
loss = criterion(outputs, labels)
|
||||
|
||||
scaler.scale(loss).backward()
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
|
||||
_, preds = torch.max(outputs, 1)
|
||||
y_true_train.extend(labels.cpu().numpy())
|
||||
y_pred_train.extend(preds.cpu().numpy())
|
||||
acc = accuracy_score(labels.cpu(), preds.cpu())
|
||||
train_loss.append(loss.item())
|
||||
train_acc.append(acc)
|
||||
|
||||
# 验证集评估
|
||||
model.eval()
|
||||
val_loss = []
|
||||
with torch.no_grad():
|
||||
for inputs, labels in val_loader:
|
||||
inputs, labels = inputs.to(device), labels.to(device)
|
||||
outputs = model(inputs)
|
||||
loss = criterion(outputs, labels)
|
||||
val_loss.append(loss.item())
|
||||
|
||||
avg_val_loss = np.mean(val_loss)
|
||||
avg_train_loss = np.mean(train_loss)
|
||||
avg_train_acc = np.mean(train_acc)
|
||||
|
||||
print(f"Epoch [{epoch+1}/{n_epochs}] - Train Loss: {avg_train_loss:.4f}, Train Acc: {avg_train_acc:.4f}, Val Loss: {avg_val_loss:.4f}")
|
||||
|
||||
if avg_val_loss < best_val_loss:
|
||||
best_val_loss = avg_val_loss
|
||||
early_stop_counter = 0
|
||||
torch.save(model.state_dict(), model_path)
|
||||
print("Model improved and saved.")
|
||||
else:
|
||||
early_stop_counter += 1
|
||||
print(f"No improvement. Early stop counter: {early_stop_counter}/{patience}")
|
||||
|
||||
if early_stop_counter >= patience:
|
||||
print("Early stopping triggered.")
|
||||
break
|
||||
|
||||
|
||||
# 训练集指标
|
||||
train_accuracy = accuracy_score(y_true_train, y_pred_train)
|
||||
train_precision = precision_score(y_true_train, y_pred_train, average='weighted')
|
||||
train_recall = recall_score(y_true_train, y_pred_train, average='weighted')
|
||||
train_f1 = f1_score(y_true_train, y_pred_train, average='weighted')
|
||||
train_cm = confusion_matrix(y_true_train, y_pred_train)
|
||||
|
||||
train_metrics = {
|
||||
"accuracy": train_accuracy,
|
||||
"precision": train_precision,
|
||||
"recall": train_recall,
|
||||
"f1_score": train_f1,
|
||||
"confusion_matrix": train_cm
|
||||
}
|
||||
|
||||
return model, train_metrics
|
||||
|
||||
# 测试函数
|
||||
def TransformerTest(X_test, y_test, BATCH_SIZE, nls, model_path):
|
||||
data_test = ZspProcess(X_test, X_test, y_test, y_test, need=True)[1]
|
||||
test_loader = DataLoader(data_test, batch_size=BATCH_SIZE, shuffle=False)
|
||||
|
||||
model = CNNWithTransformer(nls=nls).to(device)
|
||||
model.load_state_dict(torch.load(model_path))
|
||||
model.eval()
|
||||
|
||||
y_true, y_pred = [], []
|
||||
test_loss = []
|
||||
|
||||
criterion = FocalLoss(alpha=1, gamma=2).to(device) # 使用 FocalLoss
|
||||
|
||||
with torch.no_grad():
|
||||
for inputs, labels in test_loader:
|
||||
inputs, labels = inputs.to(device), labels.to(device)
|
||||
outputs = model(inputs)
|
||||
loss = criterion(outputs, labels)
|
||||
|
||||
_, preds = torch.max(outputs, 1)
|
||||
y_true.extend(labels.cpu().numpy())
|
||||
y_pred.extend(preds.cpu().numpy())
|
||||
test_loss.append(loss.item())
|
||||
|
||||
# 测试集指标
|
||||
test_accuracy = accuracy_score(y_true, y_pred)
|
||||
test_precision = precision_score(y_true, y_pred, average='weighted')
|
||||
test_recall = recall_score(y_true, y_pred, average='weighted')
|
||||
test_f1 = f1_score(y_true, y_pred, average='weighted')
|
||||
test_cm = confusion_matrix(y_true, y_pred)
|
||||
|
||||
test_metrics = {
|
||||
"accuracy": test_accuracy,
|
||||
"precision": test_precision,
|
||||
"recall": test_recall,
|
||||
"f1_score": test_f1,
|
||||
"confusion_matrix": test_cm
|
||||
}
|
||||
|
||||
print(f"Accuracy: {test_accuracy:.4f}, Precision: {test_precision:.4f}, Recall: {test_recall:.4f}, F1 Score: {test_f1:.4f}")
|
||||
print(f"Confusion Matrix:\n{test_cm}")
|
||||
return test_metrics
|
||||
|
||||
|
||||
def TransformerTrainAndTest(X,X_test, y, y_test, BATCH_SIZE, n_epochs, nls, model_path, val_split=0.2, patience=10):
|
||||
# 从训练集中划分验证集
|
||||
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=val_split, random_state=42)
|
||||
|
||||
# 训练模型并获取训练指标
|
||||
model, train_metrics = TransformerTrain(X_train, X_val, y_train, y_val, BATCH_SIZE, n_epochs, nls, model_path, patience)
|
||||
|
||||
# 测试模型并获取测试指标
|
||||
test_metrics = TransformerTest(X_test, y_test, BATCH_SIZE, nls, model_path)
|
||||
|
||||
return train_metrics, test_metrics
|
||||
Reference in New Issue
Block a user