造车先做三蹦子-之三:自制数据集(6x6数据集)230103

6*6的数据集制造、与识别：

#6*6的数据集的制作、与识别、测试、输出等


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

# 定义模型
class NeuralNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, num_classes)
        self.relu = nn.ReLU()

    def forward(self, x):
        out = self.relu(self.fc1(x))
        out = self.relu(self.fc2(out))
        out = self.fc3(out)
        return out

# 数据准备
train_data = torch.tensor([
    [[0,0,1,0,0,0],[0,0,1,0,0,0],[0,0,1,0,0,0],[0,0,1,0,0,0],[0,0,1,0,0,0],[0,0,1,0,0,0]], 
    [[0,0,0,0,0,0],[0,0,0,0,0,0],[1,1,1,1,1,1],[1,1,1,1,1,1],[0,0,0,0,0,0],[0,0,0,0,0,0]],
    # ... 其他训练数据
]  , dtype=torch.float32 )
train_labels = torch.tensor([
    [1,0,0,0,0,0,0],