82 lines
3.3 KiB
Python
82 lines
3.3 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
from sklearn.preprocessing import LabelEncoder
|
||
from torch.utils.data import DataLoader, TensorDataset
|
||
|
||
from Qtorch.Models.Qnn import Qnn
|
||
|
||
class QCNN(Qnn):
|
||
def __init__(self, X_train, y_train, X_test, y_test, labels=None, dropout_rate=0.3):
|
||
super(QCNN, self).__init__()
|
||
|
||
self.LABEL_ENCODER = LabelEncoder()
|
||
|
||
self.X_train, self.y_train, self.X_test, self.y_test = X_train, y_train, X_test, y_test
|
||
self.labels = labels
|
||
|
||
input_size = X_train.shape[1] # 输入的长度
|
||
num_classes = len(set(y_train)) # 分类数
|
||
|
||
# 网络层：卷积层 + 池化层 + 全连接层
|
||
self.layers = nn.ModuleList()
|
||
self.layers.append(nn.Conv1d(in_channels=1, out_channels=16, kernel_size=3)) # 卷积层
|
||
self.layers.append(nn.MaxPool1d(kernel_size=2)) # 池化层
|
||
self.layers.append(nn.Conv1d(in_channels=16, out_channels=32, kernel_size=3)) # 卷积层
|
||
self.layers.append(nn.MaxPool1d(kernel_size=2)) # 池化层
|
||
|
||
# 计算展平后的大小
|
||
conv_output_size = self._get_conv_output_size(input_size) # 卷积后的输出大小
|
||
print(f"Conv output size: {conv_output_size}") # 打印卷积后的输出大小
|
||
self.layers.append(nn.Linear(conv_output_size, 128)) # 全连接层
|
||
self.layers.append(nn.Linear(128, num_classes)) # 输出层
|
||
|
||
self.__init_weights()
|
||
|
||
def _get_conv_output_size(self, input_size):
|
||
# 计算卷积后的输出尺寸
|
||
x = torch.randn(1, 1, input_size) # 创建一个假的输入张量
|
||
for layer in self.layers:
|
||
x = layer(x) # 通过每一层
|
||
return int(x.numel()) # 返回展平后的输出大小
|
||
|
||
|
||
def forward(self, x):
|
||
# 通过卷积和池化层
|
||
for layer in self.layers[:-2]: # 除去最后两个 Linear 层
|
||
x = layer(x)
|
||
|
||
# 展平卷积后的输出
|
||
x = x.view(x.size(0), -1) # 这样 x 会变成 (batch_size, conv_output_size)
|
||
|
||
# 通过全连接层
|
||
x = self.layers[-2](x)
|
||
x = self.layers[-1](x)
|
||
|
||
return x
|
||
|
||
def __init_weights(self):
|
||
for m in self.modules():
|
||
if isinstance(m, nn.Conv1d) or isinstance(m, nn.Linear):
|
||
nn.init.xavier_uniform_(m.weight)
|
||
if m.bias is not None:
|
||
m.bias.data.fill_(0.01)
|
||
|
||
def _prepare_data(self):
|
||
|
||
# 将data转换为tensor形式, unsqueeze可以创建多一维度的，给
|
||
X_train_tensor = torch.tensor(self.X_train, dtype=torch.float32).unsqueeze(1)
|
||
self.y_train = self.LABEL_ENCODER.fit_transform(self.y_train)
|
||
y_train_tensor = torch.tensor(self.y_train, dtype=torch.long)
|
||
|
||
X_test_tensor = torch.tensor(self.X_test, dtype=torch.float32).unsqueeze(1)
|
||
self.y_test = self.LABEL_ENCODER.transform(self.y_test)
|
||
y_test_tensor = torch.tensor(self.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)
|
||
|
||
return train_loader, test_loader
|