帮我用pytorch写cnn代码，用于拟合两组数据，要有测试后的图形输出

好的，以下是一个简单的CNN代码，用于拟合两组数据，并且能够输出测试后的图形：

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt

# 定义CNN网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(32 * 7 * 7, 64)
        self.fc2 = nn.Linear(64, 2)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.pool(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.pool(x)
        x = x.view(-1, 32 * 7 * 7)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 定义训练和测试函数
def train(model, device, train_loader, optimizer, criterion):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += nn.functional.cross_entropy(output, target, reduction='sum').item() 
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))
    
    return test_loss

# 定义主函数
def main():
    # 加载数据
    train_data = np.load('train_data.npy')
    train_labels = np.load('train_labels.npy')
    test_data = np.load('test_data.npy')
    test_labels = np.load('test_labels.npy')

    # 转换成tensor格式
    train_data_tensor = torch.from_numpy(train_data).float().unsqueeze(1)
    train_labels_tensor = torch.from_numpy(train_labels).long()
    test_data_tensor = torch.from_numpy(test_data).float().unsqueeze(1)
    test_labels_tensor = torch.from_numpy(test_labels).long()

    # 定义超参数
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    lr = 0.001
    momentum = 0.9
    epochs = 10
    batch_size = 32

    # 初始化模型、损失函数和优化器
    model = Net().to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)

    # 训练和测试模型
    train_loss_list = []
    test_loss_list = []
    for epoch in range(1, epochs + 1):
        train(model, device, train_loader, optimizer, criterion)
        train_loss = test(model, device, train_loader)
        test_loss = test(model, device, test_loader)
        train_loss_list.append(train_loss)
        test_loss_list.append(test_loss)

    # 绘制损失图像
    plt.plot(range(1, epochs + 1), train_loss_list, label='Train Loss')
    plt.plot(range(1, epochs + 1), test_loss_list, label='Test Loss')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()
    plt.show()

    # 绘制测试输出图像
    test_data_tensor = test_data_tensor.to(device)
    output = model(test_data_tensor)
    output = output.cpu().detach().numpy()
    plt.scatter(test_data[:, 0], test_data[:, 1], c=test_labels, cmap='cool')
    plt.scatter(test_data[:, 0], test_data[:, 1], c=output.argmax(axis=1), cmap='cool', marker='x')
    plt.show()

if __name__ == '__main__':
    main()

其中，train_data.npy、train_labels.npy、test_data.npy 和 test_labels.npy 分别是训练数据、训练标签、测试数据和测试标签的 numpy 数组。你需要将这些数据准备好并保存成 npy 文件，然后在代码中加载它们。

这个代码是用于拟合两组数据，你需要将训练数据和测试数据替换成你自己的数据。同时，你也可以根据需要修改超参数和网络结构。