106 lines
2.6 KiB
Python
106 lines
2.6 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import pandas as pd
|
||
from abc import ABC, abstractmethod
|
||
|
||
from sklearn.metrics import confusion_matrix as cm
|
||
|
||
from torch.utils.data import DataLoader, TensorDataset
|
||
from Qfunctions.divSet import divSet as ds
|
||
from Qfunctions.saveToxlsx import save_to_xlsx as stx
|
||
|
||
|
||
class Qnn(nn.Module, ABC):
|
||
|
||
def __init__(self, data, labels=None, test_size=0.2, random_state=None):
|
||
super(Qnn, self).__init__()
|
||
|
||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
|
||
# 保存原始labe， 混淆矩阵使用
|
||
self.original_labels = labels
|
||
|
||
# 划分训练集和测试集
|
||
X_train, X_test, y_train, y_test, self.labels = ds(
|
||
data=data,
|
||
labels=labels,
|
||
test_size=test_size,
|
||
random_state=random_state
|
||
)
|
||
|
||
self.train_loader, self.test_loader = self.__prepare_data(
|
||
X_train=X_train,
|
||
y_train=y_train,
|
||
X_test=X_test,
|
||
y_test=y_test
|
||
)
|
||
|
||
# 定义结果
|
||
self.result = {
|
||
'acc_and_loss' : {
|
||
'epoch' : [],
|
||
'loss': [],
|
||
'train_accuracy': [],
|
||
'test_accuracy': [],
|
||
},
|
||
'confusion_matrix': None,
|
||
}
|
||
|
||
def accuracy(self, output, target):
|
||
pass
|
||
|
||
# 定义损失函数
|
||
def hinge_loss(self, output, target):
|
||
pass
|
||
|
||
@abstractmethod
|
||
def train_model(self, train_loader, test_loader, epochs):
|
||
pass
|
||
|
||
def confusion_matrix(self, test_outputs):
|
||
predicted = torch.argmax(test_outputs, dim=1)
|
||
true_label = torch.argmax(self.y_test, dim=1)
|
||
return cm(predicted.cpu(), true_label.cpu())
|
||
|
||
def fit(self, epochs = 100):
|
||
self.train_model(epochs)
|
||
|
||
def save(self, project_name):
|
||
for filename, data in self.result.items():
|
||
if filename == 'confusion_matrix':
|
||
data = pd.DataFrame(data, columns=self.original_labels, index=self.original_labels)
|
||
stx(project_name, filename, data)
|
||
else:
|
||
data = pd.DataFrame(data)
|
||
stx(project_name, filename, data)
|
||
|
||
def __prepare_data(self, X_train, y_train, X_test, y_test):
|
||
|
||
X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
|
||
y_train_tensor = torch.tensor(y_train, dtype=torch.long)
|
||
|
||
X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
|
||
y_test_tensor = torch.tensor(y_test, dtype=torch.long)
|
||
|
||
train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
|
||
test_dataset = TensorDataset(X_test_tensor, y_test_tensor)
|
||
|
||
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
|
||
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
|
||
|
||
print(train_loader, test_loader)
|
||
|
||
return train_loader, test_loader
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|