Views
No views yet
1# 定义模型
2class ClassConditionalUnet(nn.Module):
3 def __init__(self, num_classes=10, class_emb_size=4):
4 super().__init__()
5 # 将数字所属的类别映射到一个长度为class_emb_size的特征向量
6 self.class_emb = nn.Embedding(num_classes, class_emb_size)
7 # self.model就是一个不带条件的unet模型，在这里给他添加额外的输入通道，用于接收条件信息
8 self.model = UNet2DModel(
9 sample_size=28, #生成的图像是28*28
10 in_channels=1 + class_emb_size, #加入额外的输入通道
11 out_channels=1, # 输入单通道黑白数字图
12 layers_per_block=2, # 设置一个unet模块有多少个残差连接层
13 block_out_channels=(32, 64, 64),
14 down_block_types=(
15 "DownBlock2D", #普通的ResNet下采样模块
16 "AttnDownBlock2D", #含有spatial self-attention的下采样和模块
17 "AttnDownBlock2D",
18
19 ),
20 up_block_types=(
21 "AttnUpBlock2D", #含有spatial self-attention的ResNet上采样模块
22 "AttnUpBlock2D",
23 "UpBlock2D",
24 ),
25 )
26
27 def forward(self, x, t, class_labels):
28 bs, ch, w, h = x.shape
29 class_cond = self.class_emb(class_labels) # 将类别映射为向量形式
30 class_cond = class_cond.view(bs, class_cond.shape[1], 1, 1).expand(bs, class_cond.shape[1], w,
31 h) # 拓展张量形状
32 net_input = torch.cat((class_cond, x), dim=1)
33 return self.model(net_input, t).sample
34 ckpt = torch.load("class_cond_unet.pth", map_location="cpu")
35
36 model = ClassConditionalUnet(
37 num_classes=ckpt["num_classes"],
38 class_emb_size=ckpt["class_emb_size"]
39)
40
41 model.load_state_dict(ckpt["model_state_dict"])
42 model.eval()
43 noise_scheduler = DDPMScheduler.from_pretrained("Dirry525/class_num_generator")