如何使用pytorch实现LocallyConnected1D

import torch
import torch.nn as nn
class LocallyConnected1D(nn.Module):
    def __init__(self, input_channels, output_channels, output_length, kernel_size):
        super(LocallyConnected1D, self).__init__()
        self.output_length = output_length
        self.kernel_size = kernel_size
        # Weight tensor
        self.weight = nn.Parameter(torch.randn(output_length, input_channels, kernel_size, output_channels))
        self.bias = nn.Parameter(torch.randn(output_length, output_channels))
    def forward(self, x):
        outputs = []
        for i in range(self.output_length):
            local_input = x[:, :, i:i+self.kernel_size]
            local_output = (local_input.unsqueeze(-1) * self.weight[i]).sum(dim=2) + self.bias[i]
            outputs.append(local_output)
        return torch.stack(outputs, dim=2)

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, random_split, TensorDataset
# Generate random data
n_samples = 960
input_size = 33
X = torch.randn(n_samples, 1, input_size)
y = torch.randint(0, 2, (n_samples,))
# Split into train and validation sets
dataset = TensorDataset(X, y)
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
# Define model
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.lc = LocallyConnected1D(1, 16, 29, 5)
        self.fc = nn.Linear(29*16, 2)
    def forward(self, x):
        x = self.lc(x)
        x = x.view(x.size(0), -1)
        return self.fc(x)
model = Model()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# Training and validation
num_epochs = 10
for epoch in range(num_epochs):
    # Training
    model.train()
    train_loss = 0
    for batch_x, batch_y in train_loader:
        optimizer.zero_grad()
        outputs = model(batch_x)
        loss = criterion(outputs, batch_y)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
    # Validation
    model.eval()
    val_loss = 0
    with torch.no_grad():
        for batch_x, batch_y in val_loader:
            outputs = model(batch_x)
            loss = criterion(outputs, batch_y)
            val_loss += loss.item()
    print(f"Epoch {epoch + 1}/{num_epochs}, "
          f"Training Loss: {train_loss / len(train_loader)}, "
          f"Validation Loss: {val_loss / len(val_loader)}")