(self, opt)
| 25 | |
| 26 | class SIEN_Model(BaseModel): |
| 27 | def __init__(self, opt): |
| 28 | super(SIEN_Model, self).__init__(opt) |
| 29 | |
| 30 | self.rank = -1 # non dist training |
| 31 | train_opt = opt['train'] |
| 32 | |
| 33 | # define network and load pretrained models |
| 34 | self.netG = networks.define_G(opt).to(self.device) |
| 35 | if opt['dist']: |
| 36 | self.netG = DistributedDataParallel(self.netG, device_ids=[torch.cuda.current_device()]) |
| 37 | else: |
| 38 | self.netG = DataParallel(self.netG) |
| 39 | # print network |
| 40 | self.print_network() |
| 41 | self.load() |
| 42 | |
| 43 | ####################### Continue learning model parameter setting |
| 44 | if train_opt['ewc']: |
| 45 | self.Importance_Pre = torch.load(os.path.join(self.opt['path']['pretrain'], 'Importance.pth')) |
| 46 | self.Star_vals_Pre = torch.load(os.path.join(self.opt['path']['pretrain'], 'Star.pth')) |
| 47 | logger.info("Load Pretrain Importance and Stars!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!") |
| 48 | else: |
| 49 | self.Importance = [] |
| 50 | self.Star_vals = [] |
| 51 | for w in self.netG.parameters(): |
| 52 | self.Importance.append(torch.zeros_like(w)) |
| 53 | self.Star_vals.append(torch.zeros_like(w)) |
| 54 | logger.info("Initial Importance and Stars with zeros!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!") |
| 55 | |
| 56 | ############################### Distilation setting |
| 57 | if train_opt['distill']: |
| 58 | self.netG_Pre = networks.define_G(opt).to(self.device) |
| 59 | self.netG_Pre = DataParallel(self.netG_Pre) |
| 60 | self.load_Pre() |
| 61 | self.netG_Pre.eval() |
| 62 | #################################################################### |
| 63 | if self.is_train: |
| 64 | self.netG.train() |
| 65 | |
| 66 | #### loss |
| 67 | loss_type = train_opt['pixel_criterion'] |
| 68 | if loss_type == 'l1': |
| 69 | self.cri_pix = nn.L1Loss().to(self.device) |
| 70 | self.cri_ssim = SSIMLoss().to(self.device) |
| 71 | self.mse = nn.MSELoss().to(self.device) |
| 72 | # self.cri_vgg = VGGLoss(id=4).to(self.device) |
| 73 | self.bce = nn.BCEWithLogitsLoss().to(self.device) |
| 74 | elif loss_type == 'l2': |
| 75 | self.cri_pix = nn.MSELoss().to(self.device) |
| 76 | self.cri_ssim = SSIMLoss().to(self.device) |
| 77 | elif loss_type == 'cb': |
| 78 | self.cri_pix = CharbonnierLoss().to(self.device) |
| 79 | self.cri_ssim = SSIMLoss().to(self.device) |
| 80 | # self.cri_vgg = VGGLoss(id=4).to(self.device) |
| 81 | else: |
| 82 | raise NotImplementedError('Loss type [{:s}] is not recognized.'.format(loss_type)) |
| 83 | |
| 84 |
nothing calls this directly
no test coverage detected