搭建神经网络

1class MyUNet(nn.Module):
2    def __init__(self, n_steps=1000, time_emb_dim=100):
3        super(MyUNet, self).__init__()
4
5        # 正弦位置编码嵌入
6        self.time_embed = nn.Embedding(n_steps, time_emb_dim)
7        self.time_embed.weight.data = sinusoidal_embedding(n_steps, time_emb_dim)
8        self.time_embed.requires_grad_(False)
9
10        # 前半部分（下采样路径）
11        self.te1 = self._make_te(time_emb_dim, 1)
12        self.b1 = nn.Sequential(
13            MyBlock((1, 28, 28), 1, 10),
14            MyBlock((10, 28, 28), 10, 10),
15            MyBlock((10, 28, 28), 10, 10)
16        )
17        self.down1 = nn.Conv2d(10, 10, 4, 2, 1)
18
19        self.te2 = self._make_te(time_emb_dim, 10)
20        self.b2 = nn.Sequential(
21            MyBlock((10, 14, 14), 10, 20),
22            MyBlock((20, 14, 14), 20, 20),
23            MyBlock((20, 14, 14), 20, 20)
24        )
25        self.down2 = nn.Conv2d(20, 20, 4, 2, 1)
26
27        self.te3 = self._make_te(time_emb_dim, 20)
28        self.b3 = nn.Sequential(
29            MyBlock((20, 7, 7), 20, 40),
30            MyBlock((40, 7, 7), 40, 40),
31            MyBlock((40, 7, 7), 40, 40)
32        )
33        self.down3 = nn.Sequential(
34            nn.Conv2d(40, 40, 2, 1),
35            nn.SiLU(),
36            nn.Conv2d(40, 40, 4, 2, 1)
37        )
38
39        # 瓶颈层
40        self.te_mid = self._make_te(time_emb_dim, 40)
41        self.b_mid = nn.Sequential(
42            MyBlock((40, 3, 3), 40, 20),
43            MyBlock((20, 3, 3), 20, 20),
44            MyBlock((20, 3, 3), 20, 40)
45        )
46
47        # 后半部分（上采样路径）
48        self.up1 = nn.Sequential(
49            nn.ConvTranspose2d(40, 40, 4, 2, 1),
50            nn.SiLU(),
51            nn.ConvTranspose2d(40, 40, 2, 1)
52        )
53
54        self.te4 = self._make_te(time_emb_dim, 80)
55        self.b4 = nn.Sequential(
56            MyBlock((80, 7, 7), 80, 40),
57            MyBlock((40, 7, 7), 40, 20),
58            MyBlock((20, 7, 7), 20, 20)
59        )
60
61        self.up2 = nn.ConvTranspose2d(20, 20, 4, 2, 1)
62        self.te5 = self._make_te(time_emb_dim, 40)
63        self.b5 = nn.Sequential(
64            MyBlock((40, 14, 14), 40, 20),
65            MyBlock((20, 14, 14), 20, 10),
66            MyBlock((10, 14, 14), 10, 10)
67        )
68
69        self.up3 = nn.ConvTranspose2d(10, 10, 4, 2, 1)
70        self.te_out = self._make_te(time_emb_dim, 20)
71        self.b_out = nn.Sequential(
72            MyBlock((20, 28, 28), 20, 10),
73            MyBlock((10, 28, 28), 10, 10),
74            MyBlock((10, 28, 28), 10, 10, normalize=False)
75        )
76
77        self.conv_out = nn.Conv2d(10, 1, 3, 1, 1)
78
79    def forward(self, x, t):
80        # x 的形状为 (N, 1, 28, 28)（单通道灰度图像）
81        t = self.time_embed(t)
82        n = len(x)
83        out1 = self.b1(x + self.te1(t).reshape(n, -1, 1, 1))  # (N, 10, 28, 28)
84        out2 = self.b2(self.down1(out1) + self.te2(t).reshape(n, -1, 1, 1))  # (N, 20, 14, 14)
85        out3 = self.b3(self.down2(out2) + self.te3(t).reshape(n, -1, 1, 1))  # (N, 40, 7, 7)
86
87        out_mid = self.b_mid(self.down3(out3) + self.te_mid(t).reshape(n, -1, 1, 1))  # (N, 40, 3, 3)
88
89        out4 = torch.cat((out3, self.up1(out_mid)), dim=1)  # (N, 80, 7, 7)
90        out4 = self.b4(out4 + self.te4(t).reshape(n, -1, 1, 1))  # (N, 20, 7, 7)
91
92        out5 = torch.cat((out2, self.up2(out4)), dim=1)  # (N, 40, 14, 14)
93        out5 = self.b5(out5 + self.te5(t).reshape(n, -1, 1, 1))  # (N, 10, 14, 14)
94
95        out = torch.cat((out1, self.up3(out5)), dim=1)  # (N, 20, 28, 28)
96        out = self.b_out(out + self.te_out(t).reshape(n, -1, 1, 1))  # (N, 1, 28, 28)
97
98        out = self.conv_out(out)
99
100        return out
101
102    def _make_te(self, dim_in, dim_out):
103        return nn.Sequential(
104            nn.Linear(dim_in, dim_out),
105            nn.SiLU(),
106            nn.Linear(dim_out, dim_out)
107        )