Dirry525/class_num_generator

This model is a diffusion model for generate number.

Usage

# 定义模型
class ClassConditionalUnet(nn.Module):
    def __init__(self, num_classes=10, class_emb_size=4):
        super().__init__()
        # 将数字所属的类别映射到一个长度为class_emb_size的特征向量
        self.class_emb = nn.Embedding(num_classes, class_emb_size)
        # self.model就是一个不带条件的unet模型，在这里给他添加额外的输入通道，用于接收条件信息
        self.model = UNet2DModel(
            sample_size=28,  #生成的图像是28*28
            in_channels=1 + class_emb_size,  #加入额外的输入通道
            out_channels=1,  # 输入单通道黑白数字图
            layers_per_block=2,  # 设置一个unet模块有多少个残差连接层
            block_out_channels=(32, 64, 64),
            down_block_types=(
                "DownBlock2D",  #普通的ResNet下采样模块
                "AttnDownBlock2D",  #含有spatial self-attention的下采样和模块
                "AttnDownBlock2D",

            ),
            up_block_types=(
                "AttnUpBlock2D",  #含有spatial self-attention的ResNet上采样模块
                "AttnUpBlock2D",
                "UpBlock2D",
            ),
        )

    def forward(self, x, t, class_labels):
        bs, ch, w, h = x.shape
        class_cond = self.class_emb(class_labels)  # 将类别映射为向量形式
        class_cond = class_cond.view(bs, class_cond.shape[1], 1, 1).expand(bs, class_cond.shape[1], w,
                                                                                 h)  # 拓展张量形状
        net_input = torch.cat((class_cond, x), dim=1)
        return self.model(net_input, t).sample        
    ckpt = torch.load("class_cond_unet.pth", map_location="cpu")

    model = ClassConditionalUnet(
    num_classes=ckpt["num_classes"],
    class_emb_size=ckpt["class_emb_size"]
)

    model.load_state_dict(ckpt["model_state_dict"])
    model.eval()
    noise_scheduler = DDPMScheduler.from_pretrained("Dirry525/class_num_generator")