单隐层前馈网络求最小值

import torch
import torch.nn as nn


def min_loss(x, y):
    return (x-y)**2


class funcnet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(4, 20)
        self.relu1 = torch.tanh
        self.fc2 = nn.Linear(20, 1)
        # self.relu2 = nn.LeakyReLU(negative_slope=0.1)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu1(x)
        x = self.fc2(x)
        # x = self.relu2(x)
        return x


model = funcnet()
model.load_state_dict(torch.load('minnet.pth'))
pred = model(torch.tensor([10., 2.,3., 15]))
optimizer = torch.optim.Adam(model.parameters(), 0.001)
for epoch in range(100):
    for n in range(1000):
        data = 10*abs(torch.randn(4))
        label = torch.min(data)
        pred = model(data)
        train_loss = min_loss(pred, label)
        optimizer.zero_grad()
        train_loss.backward()
        optimizer.step()
    print(f'loss:{train_loss.item()}')
torch.save(model.state_dict(), 'minnet.pth')