61 lines
2.3 KiB
Python
61 lines
2.3 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from sklearn.preprocessing import LabelEncoder
|
|
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)
|