devbox
zyabc
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

深入浅出 diffusion(4)：pytorch 实现简单 diffusion

作者：zyabc | 2024-01-30 15:56:58

赞

踩

深入浅出 diffusion(4)：pytorch 实现简单 diffusion

 1. 训练和采样流程

 2. 无条件实现

  1. import torch, time, os
  2. import numpy as np
  3. import torch.nn as nn
  4. import torch.optim as optim
  5. from torchvision.datasets import MNIST
  6. from torchvision import transforms
  7. from torch.utils.data import DataLoader
  8. from torchvision.utils import save_image
  9. import torch.nn.functional as F
  10.  
  11.  
  12. class ResidualConvBlock(nn.Module):
  13.     def __init__(
  14.         self, in_channels: int, out_channels: int, is_res: bool = False
  15.     ) -> None:
  16.         super().__init__()
  17.         '''
  18.         standard ResNet style convolutional block
  19.         '''
  20.         self.same_channels = in_channels==out_channels
  21.         self.is_res = is_res
  22.         self.conv1 = nn.Sequential(
  23.             nn.Conv2d(in_channels, out_channels, 3, 1, 1),
  24.             nn.BatchNorm2d(out_channels),
  25.             nn.GELU(),
  26.         )
  27.         self.conv2 = nn.Sequential(
  28.             nn.Conv2d(out_channels, out_channels, 3, 1, 1),
  29.             nn.BatchNorm2d(out_channels),
  30.             nn.GELU(),
  31.         )
  32.  
  33.     def forward(self, x: torch.Tensor) -> torch.Tensor:
  34.         if self.is_res:
  35.             x1 = self.conv1(x)
  36.             x2 = self.conv2(x1)
  37.             # this adds on correct residual in case channels have increased
  38.             if self.same_channels:
  39.                 out = x + x2
  40.             else:
  41.                 out = x1 + x2
  42.             return out / 1.414
  43.         else:
  44.             x1 = self.conv1(x)
  45.             x2 = self.conv2(x1)
  46.             return x2
  47.  
  48.  
  49. class UnetDown(nn.Module):
  50.     def __init__(self, in_channels, out_channels):
  51.         super(UnetDown, self).__init__()
  52.         '''
  53.         process and downscale the image feature maps
  54.         '''
  55.         layers = [ResidualConvBlock(in_channels, out_channels), nn.MaxPool2d(2)]
  56.         self.model = nn.Sequential(*layers)
  57.  
  58.     def forward(self, x):
  59.         return self.model(x)
  60.  
  61.  
  62. class UnetUp(nn.Module):
  63.     def __init__(self, in_channels, out_channels):
  64.         super(UnetUp, self).__init__()
  65.         '''
  66.         process and upscale the image feature maps
  67.         '''
  68.         layers = [
  69.             nn.ConvTranspose2d(in_channels, out_channels, 2, 2),
  70.             ResidualConvBlock(out_channels, out_channels),
  71.             ResidualConvBlock(out_channels, out_channels),
  72.         ]
  73.         self.model = nn.Sequential(*layers)
  74.  
  75.     def forward(self, x, skip):
  76.         x = torch.cat((x, skip), 1)
  77.         x = self.model(x)
  78.         return x
  79.  
  80.  
  81. class EmbedFC(nn.Module):
  82.     def __init__(self, input_dim, emb_dim):
  83.         super(EmbedFC, self).__init__()
  84.         '''
  85.         generic one layer FC NN for embedding things  
  86.         '''
  87.         self.input_dim = input_dim
  88.         layers = [
  89.             nn.Linear(input_dim, emb_dim),
  90.             nn.GELU(),
  91.             nn.Linear(emb_dim, emb_dim),
  92.         ]
  93.         self.model = nn.Sequential(*layers)
  94.  
  95.     def forward(self, x):
  96.         x = x.view(-1, self.input_dim)
  97.         return self.model(x)
  98. class Unet(nn.Module):
  99.     def __init__(self, in_channels, n_feat=256):
  100.         super(Unet, self).__init__()
  101.  
  102.         self.in_channels = in_channels
  103.         self.n_feat = n_feat
  104.  
  105.         self.init_conv = ResidualConvBlock(in_channels, n_feat, is_res=True)
  106.  
  107.         self.down1 = UnetDown(n_feat, n_feat)
  108.         self.down2 = UnetDown(n_feat, 2 * n_feat)
  109.  
  110.         self.to_vec = nn.Sequential(nn.AvgPool2d(7), nn.GELU())
  111.  
  112.         self.timeembed1 = EmbedFC(1, 2 * n_feat)
  113.         self.timeembed2 = EmbedFC(1, 1 * n_feat)
  114.  
  115.         self.up0 = nn.Sequential(
  116.             # nn.ConvTranspose2d(6 * n_feat, 2 * n_feat, 7, 7), # when concat temb and cemb end up w 6*n_feat
  117.             nn.ConvTranspose2d(2 * n_feat, 2 * n_feat, 7, 7),  # otherwise just have 2*n_feat
  118.             nn.GroupNorm(8, 2 * n_feat),
  119.             nn.ReLU(),
  120.         )
  121.  
  122.         self.up1 = UnetUp(4 * n_feat, n_feat)
  123.         self.up2 = UnetUp(2 * n_feat, n_feat)
  124.         self.out = nn.Sequential(
  125.             nn.Conv2d(2 * n_feat, n_feat, 3, 1, 1),
  126.             nn.GroupNorm(8, n_feat),
  127.             nn.ReLU(),
  128.             nn.Conv2d(n_feat, self.in_channels, 3, 1, 1),
  129.         )
  130.  
  131.     def forward(self, x, t):
  132.         '''
  133.         输入加噪图像和对应的时间step，预测反向噪声的正态分布
  134.         :param x: 加噪图像
  135.         :param t: 对应step
  136.         :return: 正态分布噪声
  137.         '''
  138.         x = self.init_conv(x)
  139.         down1 = self.down1(x)
  140.         down2 = self.down2(down1)
  141.         hiddenvec = self.to_vec(down2)
  142.  
  143.         # embed time step
  144.         temb1 = self.timeembed1(t).view(-1, self.n_feat * 2, 1, 1)
  145.         temb2 = self.timeembed2(t).view(-1, self.n_feat, 1, 1)
  146.  
  147.         # 将上采样输出与step编码相加，输入到下一个上采样层
  148.         up1 = self.up0(hiddenvec)
  149.         up2 = self.up1(up1 + temb1, down2)
  150.         up3 = self.up2(up2 + temb2, down1)
  151.         out = self.out(torch.cat((up3, x), 1))
  152.         return out
  153.  
  154. class DDPM(nn.Module):
  155.     def __init__(self, model, betas, n_T, device):
  156.         super(DDPM, self).__init__()
  157.         self.model = model.to(device)
  158.  
  159.         # register_buffer 可以提前保存alpha相关，节约时间
  160.         for k, v in self.ddpm_schedules(betas[0], betas[1], n_T).items():
  161.             self.register_buffer(k, v)
  162.  
  163.         self.n_T = n_T
  164.         self.device = device
  165.         self.loss_mse = nn.MSELoss()
  166.  
  167.     def ddpm_schedules(self, beta1, beta2, T):
  168.         '''
  169.         提前计算各个step的alpha，这里beta是线性变化
  170.         :param beta1: beta的下限
  171.         :param beta2: beta的下限
  172.         :param T: 总共的step数
  173.         '''
  174.         assert beta1 < beta2 < 1.0, "beta1 and beta2 must be in (0, 1)"
  175.  
  176.         beta_t = (beta2 - beta1) * torch.arange(0, T + 1, dtype=torch.float32) / T + beta1 # 生成beta1-beta2均匀分布的数组
  177.         sqrt_beta_t = torch.sqrt(beta_t)
  178.         alpha_t = 1 - beta_t
  179.         log_alpha_t = torch.log(alpha_t)
  180.         alphabar_t = torch.cumsum(log_alpha_t, dim=0).exp() # alpha累乘
  181.  
  182.         sqrtab = torch.sqrt(alphabar_t) # 根号alpha累乘
  183.         oneover_sqrta = 1 / torch.sqrt(alpha_t) # 1 / 根号alpha
  184.  
  185.         sqrtmab = torch.sqrt(1 - alphabar_t) # 根号下（1-alpha累乘）
  186.         mab_over_sqrtmab_inv = (1 - alpha_t) / sqrtmab
  187.  
  188.         return {
  189.             "alpha_t": alpha_t,  # \alpha_t
  190.             "oneover_sqrta": oneover_sqrta,  # 1/\sqrt{\alpha_t}
  191.             "sqrt_beta_t": sqrt_beta_t,  # \sqrt{\beta_t}
  192.             "alphabar_t": alphabar_t,  # \bar{\alpha_t}
  193.             "sqrtab": sqrtab,  # \sqrt{\bar{\alpha_t}} # 加噪标准差
  194.             "sqrtmab": sqrtmab,  # \sqrt{1-\bar{\alpha_t}}  # 加噪均值
  195.             "mab_over_sqrtmab": mab_over_sqrtmab_inv,  # (1-\alpha_t)/\sqrt{1-\bar{\alpha_t}}
  196.         }
  197.     def forward(self, x):
  198.         """
  199.         训练过程中, 随机选择step和生成噪声
  200.         """
  201.         # 随机选择step
  202.         _ts = torch.randint(1, self.n_T + 1, (x.shape[0],)).to(self.device)  # t ~ Uniform(0, n_T)
  203.         # 随机生成正态分布噪声
  204.         noise = torch.randn_like(x)  # eps ~ N(0, 1)
  205.         # 加噪后的图像x_t
  206.         x_t = (
  207.                 self.sqrtab[_ts, None, None, None] * x
  208.                 + self.sqrtmab[_ts, None, None, None] * noise
  209.  
  210.         )
  211.  
  212.         # 将unet预测的对应step的正态分布噪声与真实噪声做对比
  213.         return self.loss_mse(noise, self.model(x_t, _ts / self.n_T))
  214.  
  215.     def sample(self, n_sample, size, device):
  216.         # 随机生成初始噪声图片 x_T ~ N(0, 1)
  217.         x_i = torch.randn(n_sample, *size).to(device)
  218.         for i in range(self.n_T, 0, -1):
  219.             t_is = torch.tensor([i / self.n_T]).to(device)
  220.             t_is = t_is.repeat(n_sample, 1, 1, 1)
  221.  
  222.             z = torch.randn(n_sample, *size).to(device) if i > 1 else 0
  223.  
  224.             eps = self.model(x_i, t_is)
  225.             x_i = x_i[:n_sample]
  226.             x_i = self.oneover_sqrta[i] * (x_i - eps * self.mab_over_sqrtmab[i]) + self.sqrt_beta_t[i] * z
  227.         return x_i
  228.  
  229.  
  230. class ImageGenerator(object):
  231.     def __init__(self):
  232.         '''
  233.         初始化，定义超参数、数据集、网络结构等
  234.         '''
  235.         self.epoch = 20
  236.         self.sample_num = 100
  237.         self.batch_size = 256
  238.         self.lr = 0.0001
  239.         self.n_T = 400
  240.         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
  241.         self.init_dataloader()
  242.         self.sampler = DDPM(model=Unet(in_channels=1), betas=(1e-4, 0.02), n_T=self.n_T, device=self.device).to(self.device)
  243.         self.optimizer = optim.Adam(self.sampler.model.parameters(), lr=self.lr)
  244.  
  245.     def init_dataloader(self):
  246.         '''
  247.         初始化数据集和dataloader
  248.         '''
  249.         tf = transforms.Compose([
  250.             transforms.ToTensor(),
  251.         ])
  252.         train_dataset = MNIST('./data/',
  253.                               train=True,
  254.                               download=True,
  255.                               transform=tf)
  256.         self.train_dataloader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True, drop_last=True)
  257.         val_dataset = MNIST('./data/',
  258.                             train=False,
  259.                             download=True,
  260.                             transform=tf)
  261.         self.val_dataloader = DataLoader(val_dataset, batch_size=self.batch_size, shuffle=False)
  262.  
  263.     def train(self):
  264.         self.sampler.train()
  265.         print('训练开始!!')
  266.         for epoch in range(self.epoch):
  267.             self.sampler.model.train()
  268.             loss_mean = 0
  269.             for i, (images, labels) in enumerate(self.train_dataloader):
  270.                 images, labels = images.to(self.device), labels.to(self.device)
  271.  
  272.                 # 将latent和condition拼接后输入网络
  273.                 loss = self.sampler(images)
  274.                 loss_mean += loss.item()
  275.                 self.optimizer.zero_grad()
  276.                 loss.backward()
  277.                 self.optimizer.step()
  278.             train_loss = loss_mean / len(self.train_dataloader)
  279.             print('epoch:{}, loss:{:.4f}'.format(epoch, train_loss))
  280.             self.visualize_results(epoch)
  281.  
  282.     @torch.no_grad()
  283.     def visualize_results(self, epoch):
  284.         self.sampler.eval()
  285.         # 保存结果路径
  286.         output_path = 'results/Diffusion'
  287.         if not os.path.exists(output_path):
  288.             os.makedirs(output_path)
  289.  
  290.         tot_num_samples = self.sample_num
  291.         image_frame_dim = int(np.floor(np.sqrt(tot_num_samples)))
  292.         out = self.sampler.sample(tot_num_samples, (1, 28, 28), self.device)
  293.         save_image(out, os.path.join(output_path, '{}.jpg'.format(epoch)), nrow=image_frame_dim)
  294.  
  295.  
  296.  
  297. if __name__ == '__main__':
  298.     generator = ImageGenerator()
  299.     generator.train()

3. 有条件实现

  1. import torch, time, os
  2. import numpy as np
  3. import torch.nn as nn
  4. import torch.optim as optim
  5. from torchvision.datasets import MNIST
  6. from torchvision import transforms
  7. from torch.utils.data import DataLoader
  8. from torchvision.utils import save_image
  9. import torch.nn.functional as F
  10.  
  11.  
  12. class ResidualConvBlock(nn.Module):
  13.     def __init__(
  14.         self, in_channels: int, out_channels: int, is_res: bool = False
  15.     ) -> None:
  16.         super().__init__()
  17.         '''
  18.         standard ResNet style convolutional block
  19.         '''
  20.         self.same_channels = in_channels==out_channels
  21.         self.is_res = is_res
  22.         self.conv1 = nn.Sequential(
  23.             nn.Conv2d(in_channels, out_channels, 3, 1, 1),
  24.             nn.BatchNorm2d(out_channels),
  25.             nn.GELU(),
  26.         )
  27.         self.conv2 = nn.Sequential(
  28.             nn.Conv2d(out_channels, out_channels, 3, 1, 1),
  29.             nn.BatchNorm2d(out_channels),
  30.             nn.GELU(),
  31.         )
  32.  
  33.     def forward(self, x: torch.Tensor) -> torch.Tensor:
  34.         if self.is_res:
  35.             x1 = self.conv1(x)
  36.             x2 = self.conv2(x1)
  37.             # this adds on correct residual in case channels have increased
  38.             if self.same_channels:
  39.                 out = x + x2
  40.             else:
  41.                 out = x1 + x2
  42.             return out / 1.414
  43.         else:
  44.             x1 = self.conv1(x)
  45.             x2 = self.conv2(x1)
  46.             return x2
  47.  
  48.  
  49. class UnetDown(nn.Module):
  50.     def __init__(self, in_channels, out_channels):
  51.         super(UnetDown, self).__init__()
  52.         '''
  53.         process and downscale the image feature maps
  54.         '''
  55.         layers = [ResidualConvBlock(in_channels, out_channels), nn.MaxPool2d(2)]
  56.         self.model = nn.Sequential(*layers)
  57.  
  58.     def forward(self, x):
  59.         return self.model(x)
  60.  
  61.  
  62. class UnetUp(nn.Module):
  63.     def __init__(self, in_channels, out_channels):
  64.         super(UnetUp, self).__init__()
  65.         '''
  66.         process and upscale the image feature maps
  67.         '''
  68.         layers = [
  69.             nn.ConvTranspose2d(in_channels, out_channels, 2, 2),
  70.             ResidualConvBlock(out_channels, out_channels),
  71.             ResidualConvBlock(out_channels, out_channels),
  72.         ]
  73.         self.model = nn.Sequential(*layers)
  74.  
  75.     def forward(self, x, skip):
  76.         x = torch.cat((x, skip), 1)
  77.         x = self.model(x)
  78.         return x
  79.  
  80.  
  81. class EmbedFC(nn.Module):
  82.     def __init__(self, input_dim, emb_dim):
  83.         super(EmbedFC, self).__init__()
  84.         '''
  85.         generic one layer FC NN for embedding things  
  86.         '''
  87.         self.input_dim = input_dim
  88.         layers = [
  89.             nn.Linear(input_dim, emb_dim),
  90.             nn.GELU(),
  91.             nn.Linear(emb_dim, emb_dim),
  92.         ]
  93.         self.model = nn.Sequential(*layers)
  94.  
  95.     def forward(self, x):
  96.         x = x.view(-1, self.input_dim)
  97.         return self.model(x)
  98. class Unet(nn.Module):
  99.     def __init__(self, in_channels, n_feat=256, n_classes=10):
  100.         super(Unet, self).__init__()
  101.  
  102.         self.in_channels = in_channels
  103.         self.n_feat = n_feat
  104.  
  105.         self.init_conv = ResidualConvBlock(in_channels, n_feat, is_res=True)
  106.  
  107.         self.down1 = UnetDown(n_feat, n_feat)
  108.         self.down2 = UnetDown(n_feat, 2 * n_feat)
  109.  
  110.         self.to_vec = nn.Sequential(nn.AvgPool2d(7), nn.GELU())
  111.  
  112.         self.timeembed1 = EmbedFC(1, 2 * n_feat)
  113.         self.timeembed2 = EmbedFC(1, 1 * n_feat)
  114.         self.conditionembed1 = EmbedFC(n_classes, 2 * n_feat)
  115.         self.conditionembed2 = EmbedFC(n_classes, 1 * n_feat)
  116.  
  117.         self.up0 = nn.Sequential(
  118.             # nn.ConvTranspose2d(6 * n_feat, 2 * n_feat, 7, 7), # when concat temb and cemb end up w 6*n_feat
  119.             nn.ConvTranspose2d(2 * n_feat, 2 * n_feat, 7, 7),  # otherwise just have 2*n_feat
  120.             nn.GroupNorm(8, 2 * n_feat),
  121.             nn.ReLU(),
  122.         )
  123.  
  124.         self.up1 = UnetUp(4 * n_feat, n_feat)
  125.         self.up2 = UnetUp(2 * n_feat, n_feat)
  126.         self.out = nn.Sequential(
  127.             nn.Conv2d(2 * n_feat, n_feat, 3, 1, 1),
  128.             nn.GroupNorm(8, n_feat),
  129.             nn.ReLU(),
  130.             nn.Conv2d(n_feat, self.in_channels, 3, 1, 1),
  131.         )
  132.  
  133.     def forward(self, x, c, t):
  134.         '''
  135.         输入加噪图像和对应的时间step，预测反向噪声的正态分布
  136.         :param x: 加噪图像
  137.         :param c: contition向量
  138.         :param t: 对应step
  139.         :return: 正态分布噪声
  140.         '''
  141.         x = self.init_conv(x)
  142.         down1 = self.down1(x)
  143.         down2 = self.down2(down1)
  144.         hiddenvec = self.to_vec(down2)
  145.  
  146.         # embed time step
  147.         temb1 = self.timeembed1(t).view(-1, self.n_feat * 2, 1, 1)
  148.         temb2 = self.timeembed2(t).view(-1, self.n_feat, 1, 1)
  149.         cemb1 = self.conditionembed1(c).view(-1, self.n_feat * 2, 1, 1)
  150.         cemb2 = self.conditionembed2(c).view(-1, self.n_feat, 1, 1)
  151.  
  152.         # 将上采样输出与step编码相加，输入到下一个上采样层
  153.         up1 = self.up0(hiddenvec)
  154.         up2 = self.up1(cemb1 * up1 + temb1, down2)
  155.         up3 = self.up2(cemb2 * up2 + temb2, down1)
  156.         out = self.out(torch.cat((up3, x), 1))
  157.         return out
  158.  
  159. class DDPM(nn.Module):
  160.     def __init__(self, model, betas, n_T, device):
  161.         super(DDPM, self).__init__()
  162.         self.model = model.to(device)
  163.  
  164.         # register_buffer 可以提前保存alpha相关，节约时间
  165.         for k, v in self.ddpm_schedules(betas[0], betas[1], n_T).items():
  166.             self.register_buffer(k, v)
  167.  
  168.         self.n_T = n_T
  169.         self.device = device
  170.         self.loss_mse = nn.MSELoss()
  171.  
  172.     def ddpm_schedules(self, beta1, beta2, T):
  173.         '''
  174.         提前计算各个step的alpha，这里beta是线性变化
  175.         :param beta1: beta的下限
  176.         :param beta2: beta的下限
  177.         :param T: 总共的step数
  178.         '''
  179.         assert beta1 < beta2 < 1.0, "beta1 and beta2 must be in (0, 1)"
  180.  
  181.         beta_t = (beta2 - beta1) * torch.arange(0, T + 1, dtype=torch.float32) / T + beta1 # 生成beta1-beta2均匀分布的数组
  182.         sqrt_beta_t = torch.sqrt(beta_t)
  183.         alpha_t = 1 - beta_t
  184.         log_alpha_t = torch.log(alpha_t)
  185.         alphabar_t = torch.cumsum(log_alpha_t, dim=0).exp() # alpha累乘
  186.  
  187.         sqrtab = torch.sqrt(alphabar_t) # 根号alpha累乘
  188.         oneover_sqrta = 1 / torch.sqrt(alpha_t) # 1 / 根号alpha
  189.  
  190.         sqrtmab = torch.sqrt(1 - alphabar_t) # 根号下（1-alpha累乘）
  191.         mab_over_sqrtmab_inv = (1 - alpha_t) / sqrtmab
  192.  
  193.         return {
  194.             "alpha_t": alpha_t,  # \alpha_t
  195.             "oneover_sqrta": oneover_sqrta,  # 1/\sqrt{\alpha_t}
  196.             "sqrt_beta_t": sqrt_beta_t,  # \sqrt{\beta_t}
  197.             "alphabar_t": alphabar_t,  # \bar{\alpha_t}
  198.             "sqrtab": sqrtab,  # \sqrt{\bar{\alpha_t}} # 加噪标准差
  199.             "sqrtmab": sqrtmab,  # \sqrt{1-\bar{\alpha_t}}  # 加噪均值
  200.             "mab_over_sqrtmab": mab_over_sqrtmab_inv,  # (1-\alpha_t)/\sqrt{1-\bar{\alpha_t}}
  201.         }
  202.  
  203.     def forward(self, x, c):
  204.         """
  205.         训练过程中, 随机选择step和生成噪声
  206.         """
  207.         # 随机选择step
  208.         _ts = torch.randint(1, self.n_T + 1, (x.shape[0],)).to(self.device)  # t ~ Uniform(0, n_T)
  209.         # 随机生成正态分布噪声
  210.         noise = torch.randn_like(x)  # eps ~ N(0, 1)
  211.         # 加噪后的图像x_t
  212.         x_t = (
  213.                 self.sqrtab[_ts, None, None, None] * x
  214.                 + self.sqrtmab[_ts, None, None, None] * noise
  215.  
  216.         )
  217.  
  218.         # 将unet预测的对应step的正态分布噪声与真实噪声做对比
  219.         return self.loss_mse(noise, self.model(x_t, c, _ts / self.n_T))
  220.  
  221.     def sample(self, n_sample, c, size, device):
  222.         # 随机生成初始噪声图片 x_T ~ N(0, 1)
  223.         x_i = torch.randn(n_sample, *size).to(device)
  224.         for i in range(self.n_T, 0, -1):
  225.             t_is = torch.tensor([i / self.n_T]).to(device)
  226.             t_is = t_is.repeat(n_sample, 1, 1, 1)
  227.  
  228.             z = torch.randn(n_sample, *size).to(device) if i > 1 else 0
  229.  
  230.             eps = self.model(x_i, c, t_is)
  231.             x_i = x_i[:n_sample]
  232.             x_i = self.oneover_sqrta[i] * (x_i - eps * self.mab_over_sqrtmab[i]) + self.sqrt_beta_t[i] * z
  233.         return x_i
  234.  
  235.  
  236. class ImageGenerator(object):
  237.     def __init__(self):
  238.         '''
  239.         初始化，定义超参数、数据集、网络结构等
  240.         '''
  241.         self.epoch = 20
  242.         self.sample_num = 100
  243.         self.batch_size = 256
  244.         self.lr = 0.0001
  245.         self.n_T = 400
  246.         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
  247.         self.init_dataloader()
  248.         self.sampler = DDPM(model=Unet(in_channels=1), betas=(1e-4, 0.02), n_T=self.n_T, device=self.device).to(self.device)
  249.         self.optimizer = optim.Adam(self.sampler.model.parameters(), lr=self.lr)
  250.  
  251.     def init_dataloader(self):
  252.         '''
  253.         初始化数据集和dataloader
  254.         '''
  255.         tf = transforms.Compose([
  256.             transforms.ToTensor(),
  257.         ])
  258.         train_dataset = MNIST('./data/',
  259.                               train=True,
  260.                               download=True,
  261.                               transform=tf)
  262.         self.train_dataloader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True, drop_last=True)
  263.         val_dataset = MNIST('./data/',
  264.                             train=False,
  265.                             download=True,
  266.                             transform=tf)
  267.         self.val_dataloader = DataLoader(val_dataset, batch_size=self.batch_size, shuffle=False)
  268.  
  269.     def train(self):
  270.         self.sampler.train()
  271.         print('训练开始!!')
  272.         for epoch in range(self.epoch):
  273.             self.sampler.model.train()
  274.             loss_mean = 0
  275.             for i, (images, labels) in enumerate(self.train_dataloader):
  276.                 images, labels = images.to(self.device), labels.to(self.device)
  277.                 labels = F.one_hot(labels, num_classes=10).float()
  278.                 # 将latent和condition拼接后输入网络
  279.                 loss = self.sampler(images, labels)
  280.                 loss_mean += loss.item()
  281.                 self.optimizer.zero_grad()
  282.                 loss.backward()
  283.                 self.optimizer.step()
  284.             train_loss = loss_mean / len(self.train_dataloader)
  285.             print('epoch:{}, loss:{:.4f}'.format(epoch, train_loss))
  286.             self.visualize_results(epoch)
  287.  
  288.     @torch.no_grad()
  289.     def visualize_results(self, epoch):
  290.         self.sampler.eval()
  291.         # 保存结果路径
  292.         output_path = 'results/Diffusion'
  293.         if not os.path.exists(output_path):
  294.             os.makedirs(output_path)
  295.  
  296.         tot_num_samples = self.sample_num
  297.         image_frame_dim = int(np.floor(np.sqrt(tot_num_samples)))
  298.         labels = F.one_hot(torch.Tensor(np.repeat(np.arange(10), 10)).to(torch.int64), num_classes=10).to(self.device).float()
  299.         out = self.sampler.sample(tot_num_samples, labels, (1, 28, 28), self.device)
  300.         save_image(out, os.path.join(output_path, '{}.jpg'.format(epoch)), nrow=image_frame_dim)
  301.  
  302.  
  303.  
  304. if __name__ == '__main__':
  305.     generator = ImageGenerator()
  306.     generator.train()

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/article/detail/47400?site
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号