模型代码

models/ddm.py

@@ -100,6 +100,7 @@ class DenoisingDiffusion(object):
         self.writer = SummaryWriter(config.data.tensorboard)
         self.model = DiffusionUNet(config)
         self.model.to(self.device)
+        self.test = test
         if test:
             self.model = torch.nn.DataParallel(self.model)
         else:
@@ -129,7 +130,8 @@ class DenoisingDiffusion(object):
         self.model.load_state_dict(checkpoint['state_dict'], strict=True)
         self.optimizer.load_state_dict(checkpoint['optimizer'])
         self.ema_helper.load_state_dict(checkpoint['ema_helper'])
-        self.scheduler.load_state_dict(checkpoint['scheduler'])
+        if not self.test:
+            self.scheduler.load_state_dict(checkpoint['scheduler'])
         if ema:
             self.ema_helper.ema(self.model)
         print("=> loaded checkpoint '{}' (epoch {}, step {})".format(load_path, checkpoint['epoch'], self.step))