File size: 2,281 Bytes
212a3ca |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
---
license: mit
tags:
- pytorch
- diffusers
- class-conditional-image-generation
- diffusion-models-class
---
# This model is a diffusion model for generate number.
## Usage
```python
# 定义模型
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")
```
|