1
import torch
2
from torch import nn, Tensor
3

4
import matplotlib.pyplot as plt
5
from sklearn.datasets import make_moons
6

7
# model
8
class Flow(nn.Module):
9
    def __init__(self, dim: int = 2, h: int = 64):
10
        super().__init__()
11
        self.net = nn.Sequential(
12
            nn.Linear(dim + 1, h), nn.ELU(),
13
            nn.Linear(h, h), nn.ELU(),
14
            nn.Linear(h, h), nn.ELU(),
15
            nn.Linear(h, dim))
16

17
    def forward(self, t: Tensor, x_t: Tensor) -> Tensor:
18
        return self.net(torch.cat((t, x_t), -1))
19

20
    def step(self, x_t: Tensor, t_start: Tensor, t_end: Tensor) -> Tensor
21
    # 这个类实现了中点法（二阶Runge-Kutta方法）来数值求解ODE：
22
      """
23
      t_start.view(1, 1): 如果 t_start = 0.5，变成 tensor([[0.5]])
24
      .expand(x_t.shape[0], 1)将张量扩展为形状 (x_t.shape[0], 1)
25

26
      example:
27
                t_start = torch.tensor(0.5)  # 标量张量
28
                x_t = torch.randn(32, 3, 64, 64)  # 批次大小32的图像数据
29
                result = t_start.view(1, 1).expand(x_t.shape[0], 1)
30
                # result.shape: torch.Size([32, 1])
31
                # result: 32个元素值都是0.5
32
      """
33
        t_start = t_start.view(1, 1).expand(x_t.shape[0], 1)
34

35
        # self()=self.forward()
36
        # 在 PyTorch 中，当你对 nn.Module 实例使用调用语法 ()，它实际上调用的是 forward 方法
37
        return x_t + (t_end - t_start) * self(t=t_start + (t_end - t_start) / 2, x_t= x_t + self(x_t=x_t, t=t_start) * (t_end - t_start) / 2)
38
# def step(self, x_t: Tensor, t_start: Tensor, t_end: Tensor) -> Tensor:
39
#     # 1. 将时间步扩展到整个批次
40
#     t_start = t_start.view(1, 1).expand(x_t.shape[0], 1)
41
#
42
#     # 2. 计算时间步长
43
#     h = t_end - t_start  # 步长
44
#
45
#     # 3. 计算k1: f(t_start, x_t)
46
#     k1 = self(t=t_start, x_t=x_t)
47
#
48
#     # 4. 计算中点状态: x_t + h/2 * k1
49
#     midpoint_x = x_t + k1 * (h / 2)
50
#     midpoint_t = t_start + h / 2
51
#
52
#     # 5. 计算k2: f(midpoint_t, midpoint_x)
53
#     k2 = self(t=midpoint_t, x_t=midpoint_x)
54
#
55
#     # 6. 更新状态: x_t + h * k2
56
#     return x_t + h * k2
57

58

59
# train
60
flow = Flow()
61

62
optimizer = torch.optim.Adam(flow.parameters(), 1e-2)
63
loss_fn = nn.MSELoss()
64

65
for _ in range(10000):
66
    """
67
    成一个包含两个"月牙"形状的二维数据集
68
    256: 生成256个数据点
69
  noise=0.05: 添加少量噪声（5%）使数据点略微分散，更接近真实数据
70
    """
71
    x_1 = Tensor(make_moons(256, noise=0.05)[0])
72

73
    # noise
74
    x_0 = torch.randn_like(x_1)
75
    # 创建一个形状为 (len(x_1), 1) 的张量
76
    """
77
    torch.rand()：均匀分布，[0, 1)
78
    torch.randn()：标准正态分布，均值0，方差1
79
    torch.randint()：整数随机数
80
    torch.rand_like()：生成与输入张量形状相同的随机数
81
    """
82
    t = torch.rand(len(x_1), 1)
83

84
    #  最优传输条件求x_t的值
85
    x_t = (1 - t) * x_0 + t * x_1
86
    # 梯度,即x_t对t的导数
87
    dx_t = x_1 - x_0
88

89

90
    optimizer.zero_grad()
91
    # 计算网络模拟的梯度与真实梯度之间的损失
92
    loss_fn(flow(t=t, x_t=x_t), dx_t).backward()
93
    optimizer.step()
94

95

96
# sample
97

98

99
x = torch.randn(300, 2)
100
n_steps = 8
101
fig, axes = plt.subplots(1, n_steps + 1, figsize=(30, 4), sharex=True, sharey=True)
102

103
# 创建一个在区间 [0, 1.0] 上均匀分布的数值序列。
104
time_steps = torch.linspace(0, 1.0, n_steps + 1)
105

106
axes[0].scatter(x.detach()[:, 0], x.detach()[:, 1], s=10)
107
axes[0].set_title(f't = {time_steps[0]:.2f}')
108
axes[0].set_xlim(-3.0, 3.0)
109
axes[0].set_ylim(-3.0, 3.0)
110

111
for i in range(n_steps):
112
    # 计算每一时刻x_t的情况
113
    x = flow.step(x_t=x, t_start=time_steps[i], t_end=time_steps[i + 1])
114
    axes[i + 1].scatter(x.detach()[:, 0], x.detach()[:, 1], s=10)
115
    axes[i + 1].set_title(f't = {time_steps[i + 1]:.2f}')
116

117
plt.tight_layout()
118
plt.show()