deepliif 1.1.11__py3-none-any.whl → 1.1.12__py3-none-any.whl

deepliif/models/DeepLIIF_model.py

@@ -1,6 +1,7 @@
 import torch
 from .base_model import BaseModel
 from . import networks
+from .networks import get_optimizer
 
 
 class DeepLIIFModel(BaseModel):
@@ -13,13 +14,12 @@ class DeepLIIFModel(BaseModel):
             opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
         """
         BaseModel.__init__(self, opt)
-
-        # weights of the modalities in generating segmentation mask
-        self.seg_weights = [0.25, 0.15, 0.25, 0.1, 0.25]
-
-        # loss weights in calculating the final loss
-        self.loss_G_weights = [0.2, 0.2, 0.2, 0.2, 0.2]
-        self.loss_D_weights = [0.2, 0.2, 0.2, 0.2, 0.2]
+        if not hasattr(opt,'net_gs'):
+            opt.net_gs = 'unet_512'
+        
+        self.seg_weights = opt.seg_weights
+        self.loss_G_weights = opt.loss_G_weights
+        self.loss_D_weights = opt.loss_D_weights
         
         if not opt.is_train:
             self.gpu_ids = [] # avoid the models being loaded as DP
@@ -31,7 +31,7 @@ class DeepLIIFModel(BaseModel):
         # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
         for i in range(1, self.opt.modalities_no + 1 + 1):
             self.loss_names.extend(['G_GAN_' + str(i), 'G_L1_' + str(i), 'D_real_' + str(i), 'D_fake_' + str(i)])
-            self.visual_names.extend(['fake_B_' + str(i), 'real_B_' + str(i)])
+            self.visual_names.extend(['fake_B_' + str(i), 'fake_B_5' + str(i), 'real_B_' + str(i)])
 
         # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
         # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
@@ -51,24 +51,31 @@ class DeepLIIFModel(BaseModel):
                 self.model_names.extend(['G5' + str(i)])
 
         # define networks (both generator and discriminator)
-        self.netG1 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
+        if isinstance(opt.netG, str):
+            opt.netG = [opt.netG] * 4
+        if isinstance(opt.net_gs, str):
+            opt.net_gs = [opt.net_gs]*5
+
+            
+        self.netG1 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG[0], opt.norm,
                                       not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, opt.padding)
-        self.netG2 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
+        self.netG2 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG[1], opt.norm,
                                       not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, opt.padding)
-        self.netG3 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
+        self.netG3 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG[2], opt.norm,
                                       not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, opt.padding)
-        self.netG4 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
+        self.netG4 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG[3], opt.norm,
                                       not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, opt.padding)
 
-        self.netG51 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, 'unet_512', opt.norm,
+        # DeepLIIF model currently uses one gs arch because there is only one explicit seg mod output
+        self.netG51 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.net_gs[0], opt.norm,
                                       not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
-        self.netG52 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, 'unet_512', opt.norm,
+        self.netG52 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.net_gs[1], opt.norm,
                                       not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
-        self.netG53 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, 'unet_512', opt.norm,
+        self.netG53 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.net_gs[2], opt.norm,
                                       not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
-        self.netG54 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, 'unet_512', opt.norm,
+        self.netG54 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.net_gs[3], opt.norm,
                                       not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
-        self.netG55 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, 'unet_512', opt.norm,
+        self.netG55 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.net_gs[4], opt.norm,
                                       not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
 
 
@@ -101,10 +108,18 @@ class DeepLIIFModel(BaseModel):
 
             # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
             params = list(self.netG1.parameters()) + list(self.netG2.parameters()) + list(self.netG3.parameters()) + list(self.netG4.parameters()) + list(self.netG51.parameters()) + list(self.netG52.parameters()) + list(self.netG53.parameters()) + list(self.netG54.parameters()) + list(self.netG55.parameters())
-            self.optimizer_G = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
+            try:
+                self.optimizer_G = get_optimizer(opt.optimizer)(params, lr=opt.lr_g, betas=(opt.beta1, 0.999))
+            except:
+                print(f'betas are not used for optimizer torch.optim.{opt.optimizer} in generators')
+                self.optimizer_G = get_optimizer(opt.optimizer)(params, lr=opt.lr_g)
 
             params = list(self.netD1.parameters()) + list(self.netD2.parameters()) + list(self.netD3.parameters()) + list(self.netD4.parameters()) + list(self.netD51.parameters()) + list(self.netD52.parameters()) + list(self.netD53.parameters()) + list(self.netD54.parameters()) + list(self.netD55.parameters())
-            self.optimizer_D = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
+            try:
+                self.optimizer_D = get_optimizer(opt.optimizer)(params, lr=opt.lr_d, betas=(opt.beta1, 0.999))
+            except:
+                print(f'betas are not used for optimizer torch.optim.{opt.optimizer} in discriminators')
+                self.optimizer_D = get_optimizer(opt.optimizer)(params, lr=opt.lr_d)
 
             self.optimizers.append(self.optimizer_G)
             self.optimizers.append(self.optimizer_D)
@@ -322,3 +337,38 @@ class DeepLIIFModel(BaseModel):
         self.optimizer_G.zero_grad()        # set G's gradients to zero
         self.backward_G()                   # calculate graidents for G
         self.optimizer_G.step()             # udpate G's weights
+    
+    def calculate_losses(self):
+        """
+        Calculate losses but do not optimize parameters. Used in validation loss calculation during training.
+        """
+        
+        self.forward()                   # compute fake images: G(A)
+        # update D
+        self.set_requires_grad(self.netD1, True)  # enable backprop for D1
+        self.set_requires_grad(self.netD2, True)  # enable backprop for D2
+        self.set_requires_grad(self.netD3, True)  # enable backprop for D3
+        self.set_requires_grad(self.netD4, True)  # enable backprop for D4
+        self.set_requires_grad(self.netD51, True)  # enable backprop for D51
+        self.set_requires_grad(self.netD52, True)  # enable backprop for D52
+        self.set_requires_grad(self.netD53, True)  # enable backprop for D53
+        self.set_requires_grad(self.netD54, True)  # enable backprop for D54
+        self.set_requires_grad(self.netD55, True)  # enable backprop for D54
+
+        self.optimizer_D.zero_grad()        # set D's gradients to zero
+        self.backward_D()                # calculate gradients for D
+
+        # update G
+        self.set_requires_grad(self.netD1, False)  # D1 requires no gradients when optimizing G1
+        self.set_requires_grad(self.netD2, False)  # D2 requires no gradients when optimizing G2
+        self.set_requires_grad(self.netD3, False)  # D3 requires no gradients when optimizing G3
+        self.set_requires_grad(self.netD4, False)  # D4 requires no gradients when optimizing G4
+        self.set_requires_grad(self.netD51, False)  # D51 requires no gradients when optimizing G51
+        self.set_requires_grad(self.netD52, False)  # D52 requires no gradients when optimizing G52
+        self.set_requires_grad(self.netD53, False)  # D53 requires no gradients when optimizing G53
+        self.set_requires_grad(self.netD54, False)  # D54 requires no gradients when optimizing G54
+        self.set_requires_grad(self.netD55, False)  # D54 requires no gradients when optimizing G54
+
+        self.optimizer_G.zero_grad()        # set G's gradients to zero
+        self.backward_G()                   # calculate graidents for G
+            

deepliif/models/SDG_model.py

@@ -1,6 +1,7 @@
 import torch
 from .base_model import BaseModel
 from . import networks
+from .networks import get_optimizer
 
 
 class SDGModel(BaseModel):
@@ -15,28 +16,34 @@ class SDGModel(BaseModel):
         BaseModel.__init__(self, opt)
 
         self.mod_gen_no = self.opt.modalities_no
+            
 
         # weights of the modalities in generating segmentation mask
-        self.seg_weights = [0, 0, 0]
-        if opt.seg_gen:
-            self.seg_weights = [0.3] * self.mod_gen_no
-            self.seg_weights[1] = 0.4
+        self.seg_weights = opt.seg_weights
 
         # self.seg_weights = opt.seg_weights
         # assert len(self.seg_weights) == self.seg_gen_no, 'The number of the segmentation weights (seg_weights) is not equal to the number of target images (modalities_no)!'
-        # print(self.seg_weights)
+
         # loss weights in calculating the final loss
-        self.loss_G_weights = [1 / self.mod_gen_no] * self.mod_gen_no
+        self.loss_G_weights = opt.loss_G_weights
         self.loss_GS_weights = [1 / self.mod_gen_no] * self.mod_gen_no
 
-        self.loss_D_weights = [1 / self.mod_gen_no] * self.mod_gen_no
+        self.loss_D_weights = opt.loss_D_weights
         self.loss_DS_weights = [1 / self.mod_gen_no] * self.mod_gen_no
 
+        # self.gpu_ids is a possibly modifed one for model initialization
+        # self.opt.gpu_ids is the original one received in the command
+        if not opt.is_train:
+            self.gpu_ids = [] # avoid the models being loaded as DP
+        else:
+            self.gpu_ids = opt.gpu_ids
+
         self.loss_names = []
         self.visual_names = ['real_A']
         # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
         for i in range(1, self.mod_gen_no + 1):
-            self.loss_names.extend(['G_GAN_' + str(i), 'G_L1_' + str(i), 'D_real_' + str(i), 'D_fake_' + str(i)])
+            self.loss_names.extend(['G_GAN_' + str(i), 'G_L1_' + str(i), 'G_VGG_'+ str(i),
+                                    'D_real_' + str(i), 'D_fake_' + str(i)])
             self.visual_names.extend(['fake_B_' + str(i), 'real_B_' + str(i)])
 
         # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
@@ -52,41 +59,48 @@ class SDGModel(BaseModel):
                 self.model_names.extend(['G_' + str(i)])
 
         # define networks (both generator and discriminator)
+        if isinstance(opt.net_g, str):
+            self.opt.net_g = [self.opt.net_g] * self.mod_gen_no
+        if isinstance(opt.net_gs, str):
+            self.opt.net_gs = [self.opt.net_gs]*self.mod_gen_no
         self.netG = [None for _ in range(self.mod_gen_no)]
+        self.netGS = [None for _ in range(self.mod_gen_no)]
         for i in range(self.mod_gen_no):
-            self.netG[i] = networks.define_G(self.opt.input_nc * self.opt.input_no, self.opt.output_nc, self.opt.ngf, self.opt.net_g, self.opt.norm,
-                                             not self.opt.no_dropout, self.opt.init_type, self.opt.init_gain, self.opt.gpu_ids, self.opt.padding)
-            print('***************************************')
-            print(self.opt.input_nc, self.opt.output_nc, self.opt.ngf, self.opt.net_g, self.opt.norm,
-                                             not self.opt.no_dropout, self.opt.init_type, self.opt.init_gain, self.opt.gpu_ids, self.opt.padding)
-            print('***************************************')
+            self.netG[i] = networks.define_G(self.opt.input_nc * self.opt.input_no, self.opt.output_nc, self.opt.ngf, self.opt.net_g[i], self.opt.norm,
+                                             not self.opt.no_dropout, self.opt.init_type, self.opt.init_gain, self.gpu_ids, self.opt.padding)
 
         if self.is_train:  # define a discriminator; conditional GANs need to take both input and output images; Therefore, #channels for D is input_nc + output_nc
             self.netD = [None for _ in range(self.mod_gen_no)]
             for i in range(self.mod_gen_no):
                 self.netD[i] = networks.define_D(self.opt.input_nc * self.opt.input_no + self.opt.output_nc, self.opt.ndf, self.opt.net_d,
                                                  self.opt.n_layers_D, self.opt.norm, self.opt.init_type, self.opt.init_gain,
-                                                 self.opt.gpu_ids)
+                                                 self.gpu_ids)
 
         if self.is_train:
             # define loss functions
             self.criterionGAN_mod = networks.GANLoss(self.opt.gan_mode).to(self.device)
             self.criterionGAN_seg = networks.GANLoss(self.opt.gan_mode_s).to(self.device)
-
             self.criterionSmoothL1 = torch.nn.SmoothL1Loss()
-
             self.criterionVGG = networks.VGGLoss().to(self.device)
 
             # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
             params = []
             for i in range(len(self.netG)):
                 params += list(self.netG[i].parameters())
-            self.optimizer_G = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
+            try:
+                self.optimizer_G = get_optimizer(opt.optimizer)(params, lr=opt.lr_g, betas=(opt.beta1, 0.999))
+            except:
+                print(f'betas are not used for optimizer torch.optim.{opt.optimizer} in generators')
+                self.optimizer_G = get_optimizer(opt.optimizer)(params, lr=opt.lr_g)
 
             params = []
             for i in range(len(self.netD)):
                 params += list(self.netD[i].parameters())
-            self.optimizer_D = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
+            try:
+                self.optimizer_D = get_optimizer(opt.optimizer)(params, lr=opt.lr_d, betas=(opt.beta1, 0.999))
+            except:
+                print(f'betas are not used for optimizer torch.optim.{opt.optimizer} in discriminators')
+                self.optimizer_D = get_optimizer(opt.optimizer)(params, lr=opt.lr_d)
 
             self.optimizers.append(self.optimizer_G)
             self.optimizers.append(self.optimizer_D)
@@ -136,7 +150,6 @@ class SDGModel(BaseModel):
             self.loss_D_real.append(self.criterionGAN_mod(pred_real[i], True))
 
         # combine losses and calculate gradients
-        # self.loss_D = (self.loss_D_fake[0] + self.loss_D_real[0]) * 0.5 * self.loss_D_weights[0]
         self.loss_D = torch.tensor(0., device=self.device)
         for i in range(0, self.mod_gen_no):
             self.loss_D += (self.loss_D_fake[i] + self.loss_D_real[i]) * 0.5 * self.loss_D_weights[i]
@@ -159,15 +172,14 @@ class SDGModel(BaseModel):
         for i in range(self.mod_gen_no):
             self.loss_G_L1.append(self.criterionSmoothL1(self.fake_B[i], self.real_B[i]) * self.opt.lambda_L1)
 
-        #self.loss_G_VGG = []
-        #for i in range(self.mod_gen_no):
-        #    self.loss_G_VGG.append(self.criterionVGG(self.fake_B[i], self.real_B[i]) * self.opt.lambda_feat)
+        self.loss_G_VGG = []
+        for i in range(self.mod_gen_no):
+            self.loss_G_VGG.append(self.criterionVGG(self.fake_B[i], self.real_B[i]) * self.opt.lambda_feat)
+        
 
-        # self.loss_G = (self.loss_G_GAN[0] + self.loss_G_L1[0]) * self.loss_G_weights[0]
         self.loss_G = torch.tensor(0., device=self.device)
         for i in range(0, self.mod_gen_no):
-            self.loss_G += (self.loss_G_GAN[i] + self.loss_G_L1[i]) * self.loss_G_weights[i]
-        #    self.loss_G += (self.loss_G_GAN[i] + self.loss_G_L1[i] + self.loss_G_VGG[i]) * self.loss_G_weights[i]
+            self.loss_G += (self.loss_G_GAN[i] + self.loss_G_L1[i] + self.loss_G_VGG[i]) * self.loss_G_weights[i]
         self.loss_G.backward()
 
     def optimize_parameters(self):
@@ -187,3 +199,22 @@ class SDGModel(BaseModel):
         self.optimizer_G.zero_grad()  # set G's gradients to zero
         self.backward_G()  # calculate graidents for G
         self.optimizer_G.step()  # udpate G's weights
+    
+    def calculate_losses(self):
+        """
+        Calculate losses but do not optimize parameters. Used in validation loss calculation during training.
+        """
+        self.forward()  # compute fake images: G(A)
+        # update D
+        for i in range(self.mod_gen_no):
+            self.set_requires_grad(self.netD[i], True)  # enable backprop for D1
+
+        self.optimizer_D.zero_grad()  # set D's gradients to zero
+        self.backward_D()  # calculate gradients for D
+
+        # update G
+        for i in range(self.mod_gen_no):
+            self.set_requires_grad(self.netD[i], False)
+
+        self.optimizer_G.zero_grad()  # set G's gradients to zero
+        self.backward_G()  # calculate graidents for G