Blame for modelADT.py

#!/usr/bin/python
'''

Defines a generic wrapper class for all the network models

Utilizes params.py to create, initiate, load and train the network.

'''
import torch
import numpy as np
from collections import OrderedDict
import scipy.io
from torch.autograd import Variable
import os
from Generators import *
from utils import *

class ModelADT():
    def name(self):
        return 'Pix2PixModel'

    def initialize(self, opt):
        
        self.opt = opt
#        self.gpu_ids = opt.gpu_ids
        gpu_ids = []
        for i in range(torch.cuda.device_count()):
            gpu_ids.append(str(i))
        print(gpu_ids)
        self.gpu_ids = gpu_ids
        self.isTrain = opt.isTrain
        self.loss_weight = opt.lambda_L
        self.reg_loss_weight = opt.lambda_reg
        self.cross_loss_weight = opt.lambda_cross
        self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) # if self.gpu_ids else torch.device('cpu')
        
        print(self.device)
        self.model_name = self.opt.name
        self.save_dir = os.path.join(self.opt.checkpoints_save_dir, self.model_name)
        self.load_dir = os.path.join(self.opt.checkpoints_load_dir, self.model_name)
        self.results_save_dir = opt.results_dir

        self.loss_names = []
        self.visual_names = []
        self.image_paths = []

        self.loss_names = ['loss_criterion']
        self.visual_names = ['output']
        

        self.net = get_model_funct(self.opt.net)(self.opt, self.gpu_ids)
        self.net = self.net.to(self.device)

        if self.isTrain:
            if self.opt.loss_type == "L2":
                self.loss_criterion = torch.nn.MSELoss()
            elif self.opt.loss_type == "L1":
                self.loss_criterion = torch.nn.L1Loss()
            elif self.opt.loss_type == "L1_sumL2":
                self.loss_criterion = torch.nn.L1Loss()
            elif self.opt.loss_type == "L2_sumL2":
                self.loss_criterion = torch.nn.MSELoss()
            elif self.opt.loss_type == "L2_sumL1":
                self.loss_criterion = torch.nn.MSELoss()
            elif self.opt.loss_type == "L2_offset_loss":
                self.loss_criterion = torch.nn.MSELoss()
            elif self.opt.loss_type == "L1_offset_loss":
                self.loss_criterion = torch.nn.L1Loss()
            elif self.opt.loss_type == "L1_sumL2_cross":
                self.loss_criterion = torch.nn.MSELoss()
            elif self.opt.loss_type == "L2_sumL2_cross":
                self.loss_criterion = torch.nn.MSELoss()
                self.cross_loss_criterion = torch.nn.NLLLoss()

            # initialize optimizers
            self.optimizers = []
            self.optimizer = torch.optim.Adam(self.net.parameters(),lr=opt.lr, betas=(opt.beta1, 0.999), weight_decay=opt.weight_decay)
            self.optimizers.append(self.optimizer)

    # load and print networks; create schedulers
    def setup(self, opt, parser=None):
        if self.isTrain:
            self.schedulers = [get_scheduler(optimizer, opt) for optimizer in self.optimizers]

        if not self.isTrain or opt.continue_train:
            self.load_networks(opt.starting_epoch_count)
            if self.isTrain and opt.continue_train:
                self.load_training_state(opt.starting_epoch_count)
        self.print_networks(opt.verbose)

    # make models eval mode during test time
    def eval(self):
        self.net.eval()

    # used in test time, wrapping `forward` in no_grad() so we don't save
    # intermediate steps for backprop
    def test(self):
        with torch.no_grad():
            self.forward()
    
    def save_outputs(self):
        if not os.path.exists(self.results_save_dir):
            os.makedirs(self.results_save_dir, exist_ok=True)
        to_save_dict = {}
        for tensor in self.visual_names:
            tensor_val = getattr(self, tensor).data.cpu().numpy()
            to_save_dict[tensor] = tensor_val
        scipy.io.savemat(self.results_save_dir+".mat", mdict=to_save_dict)
            

    # update learning rate (called once every epoch)
    def update_learning_rate(self):
        for scheduler in self.schedulers:
            scheduler.step()
        lr = self.optimizers[0].param_groups[0]['lr']
        print('learning rate = %.7f' % lr)

    # return visualization images. train.py will display these images, and save the images to a html
    def get_current_visuals(self):
        visual_ret = OrderedDict()
        for name in self.visual_names:
            if isinstance(name, str):
                visual_ret[name] = getattr(self, name)
        return visual_ret

    # return traning losses/errors. train.py will print out these errors as debugging information
    def get_current_losses(self):
        errors_ret = OrderedDict()
        for name in self.loss_names:
            if isinstance(name, str):
                # float(...) works for both scalar tensor and float number
                errors_ret[name] = float(getattr(self, 'loss_' + name))
        return errors_ret

    # save models to the disk
    def save_networks(self, epoch):
        name = self.model_name
        save_filename = '%s_net_%s.pth' % (epoch, name)
        if not os.path.exists(self.save_dir):
            os.makedirs(self.save_dir)
        save_path = os.path.join(self.save_dir, save_filename)
        net = self.net
        print(save_path)
        print('net'+name)
        if len(self.gpu_ids) > 0 and torch.cuda.is_available() and hasattr(net, 'module'):
            torch.save(net.module.cpu().state_dict(), save_path)
        else:
            torch.save(net.cpu().state_dict(), save_path)
        net.to(self.device)

        # save optimizer, scheduler, and RNG state
        if self.isTrain:
            optim_filename = '%s_optim_%s.pth' % (epoch, self.model_name)
            optim_path = os.path.join(self.save_dir, optim_filename)
            save_dict = {
                'optimizer': self.optimizer.state_dict(),
                'scheduler': self.schedulers[0].state_dict() if self.schedulers else None,
                'rng_torch': torch.random.get_rng_state(),
                'rng_numpy': np.random.get_state(),
            }
            if torch.cuda.is_available():
                save_dict['rng_cuda'] = torch.cuda.get_rng_state_all()
            torch.save(save_dict, optim_path)

    def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):
        key = keys[i]
        if i + 1 == len(keys):  # at the end, pointing to a parameter/buffer
            if module.__class__.__name__.startswith('InstanceNorm') and (key == 'running_mean' or key == 'running_var'):
                if getattr(module, key) is None:
                    state_dict.pop('.'.join(keys))
            if module.__class__.__name__.startswith('InstanceNorm') and (key == 'num_batches_tracked'):
                state_dict.pop('.'.join(keys))
        else:
            self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1)

    # load models from the disk
    def load_networks(self, epoch, load_dir=""):
        """

        epoch (int/str): epoch index / "best" / "latest"

        """
        assert isinstance(epoch,int) or epoch=="best" or epoch=="latest"
        load_filename = f'{epoch}_net_{self.model_name}.pth'

        if load_dir:
            # use given load dir
            load_path = os.path.join(load_dir, self.model_name, load_filename) 
        else:
            # use default load dir
            load_path = os.path.join(self.load_dir, load_filename)

        net = self.net
        if isinstance(net, torch.nn.DataParallel):
            net = net.module
        print('loading the model from %s' % load_path)
        # if you are using PyTorch newer than 0.4 (e.g., built from
        # GitHub source), you can remove str() on self.device
        state_dict = torch.load(load_path, weights_only=False)
        if hasattr(state_dict, '_metadata'):
            del state_dict._metadata

        # patch InstanceNorm checkpoints prior to 0.4
        print(state_dict.keys())
        for key in list(state_dict.keys()):  # need to copy keys here because we mutate in loop
            self.__patch_instance_norm_state_dict(state_dict, net, key.split('.'))
        net.load_state_dict(state_dict)
        net = net.to(self.device)

    # load optimizer and scheduler state from disk
    def load_training_state(self, epoch):
        optim_filename = f'{epoch}_optim_{self.model_name}.pth'
        load_path = os.path.join(self.load_dir, optim_filename)
        if os.path.exists(load_path):
            print(f'loading optimizer/scheduler from {load_path}')
            state = torch.load(load_path, weights_only=False)
            self.optimizer.load_state_dict(state['optimizer'])
            if state.get('scheduler') and self.schedulers:
                self.schedulers[0].load_state_dict(state['scheduler'])
            # Restore RNG states for session continuity
            if 'rng_torch' in state:
                torch.random.set_rng_state(state['rng_torch'])
            if 'rng_numpy' in state:
                np.random.set_state(state['rng_numpy'])
            if 'rng_cuda' in state and torch.cuda.is_available():
                torch.cuda.set_rng_state_all(state['rng_cuda'])
        else:
            print(f'No optimizer state found at {load_path}, starting fresh')

    # print network information
    def print_networks(self, verbose):
        print('---------- Networks initialized -------------')
        net = self.net
        name = self.model_name
        num_params = 0
        for param in net.parameters():
            num_params += param.numel()
        if verbose:
            print(net)
        print('[Network %s] Total number of parameters : %.3f M' % (name, num_params / 1e6))
        print('-----------------------------------------------')

    # set requies_grad=Fasle to avoid computation
    def set_requires_grad(self, nets, requires_grad=False):
        if not isinstance(nets, list):
            nets = [nets]
        for net in nets:
            if net is not None:
                for param in net.parameters():
                    param.requires_grad = requires_grad

    # Set the input and target data for the network to train on/evaluate against
    def set_data(self, input, target, convert=False, shuffle_channel=False):
        shape_in = input.shape
        if shuffle_channel:
            self.input = input[:,torch.randperm(shape_in[1]),:,:]
        else:
            self.input = input
        self.target = target
        if convert:
            self.input, self.target = Variable(self.input), Variable(self.target)

        self.input = self.input.to(self.device)
        self.target = self.target.to(self.device)

    # Define the forward pass to compute loss
    def forward(self):
        self.output = self.net(self.input)
        if self.opt.loss_type != "NoLoss":
            self.loss = self.loss_weight*self.loss_criterion(self.output, self.target)

            if self.opt.loss_type == "L1_sumL2" or self.opt.loss_type == "L2_sumL2":
                self.loss += self.reg_loss_weight*torch.norm(self.output).div(self.output.numel())
                self.reg_loss = self.reg_loss_weight*torch.norm(self.output).div(self.output.numel())

            if self.opt.loss_type == "L2_sumL1":
                self.loss += self.reg_loss_weight*torch.norm(self.output,p=1).div(self.output.numel())
                self.reg_loss = self.reg_loss_weight*torch.norm(self.output,p=1).div(self.output.numel())

            if self.opt.loss_type == "L1_sumL2_cross" or self.opt.loss_type == "L2_sumL2_cross":
                self.loss += self.reg_loss_weight*torch.norm(self.output).div(self.output.numel())
                self.loss += self.cross_loss_weight*self.cross_loss_criterion(self.output.flatten(start_dim=1),self.target.flatten(start_dim=1))
            
            if self.opt.loss_type == "L1_offset_loss" or self.opt.loss_type == "L2_offset_loss":
                self.loss += self.reg_loss_weight*torch.norm(self.output).div(self.output.numel())

    def backward(self):
        self.loss.backward(retain_graph=True)
    
    def optimize_parameters(self):
        self.forward()
        self.backward()  
        self.optimizer.step()
        self.optimizer.zero_grad()