造车先做三蹦子-之二:自制数据集(5x5数据集)230102

#Jupyter Notebook231001

import torch
import torch.nn as nn
import torch.optim as optim

# 定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(25, 50)
        self.fc2 = nn.Linear(50, 6)
    
    def forward(self, x):
        x = x.view(-1, 25)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 训练数据
i = [torch.tensor([[0,0,1,0,0],
                   [0,0,1,0,0],
                   [0,0,1,0,0],
                   [0,0,1,0,0],
                   [0,0,1,0,0]], dtype=torch.float32),

     torch.tensor([[0,0,0,0,0],
                   [0,0,0,0,0],
                   [1,1,1,1,1],
                   [0,0,0,0,0],
                   [0,0,0,0,0]], dtype=torch.float32),
     
     torch.tensor([[0,0,1,0,0],
                   [0,0,1,0,0],
                   [1,1,1,1,1],
                   [0,0,1,0,0],
                   [