feat: add model (e7b4c911) · Commits · 李子南 / 2022_Spring_IPOUS

OSAP_web/OSAP_flask/models/Adaptation_model/init.py

deleted100644 → 0

+0 −67

Original line number	Diff line number	Diff line
		"""This package contains modules related to objective functions, optimizations, and network architectures.

		To add a custom model class called 'dummy', you need to add a file called 'dummy_model.py' and define a subclass DummyModel inherited from BaseModel.
		You need to implement the following five functions:
		-- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).
		-- <set_input>: unpack data from dataset and apply preprocessing.
		-- <forward>: produce intermediate results.
		-- <optimize_parameters>: calculate loss, gradients, and update network weights.
		-- <modify_commandline_options>: (optionally) add model-specific options and set default options.

		In the function <__init__>, you need to define four lists:
		-- self.loss_names (str list): specify the training losses that you want to plot and save.
		-- self.model_names (str list): define networks used in our training.
		-- self.visual_names (str list): specify the images that you want to display and save.
		-- self.optimizers (optimizer list): define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an usage.

		Now you can use the model class by specifying flag '--model dummy'.
		See our template model class 'template_model.py' for more details.
		"""

		import importlib
		from Adaptation_model.base_model import BaseModel


		def find_model_using_name(model_name):
		"""Import the module "models/[model_name]_model.py".

		In the file, the class called DatasetNameModel() will
		be instantiated. It has to be a subclass of BaseModel,
		and it is case-insensitive.
		"""
		model_filename = "models." + model_name + "_model"
		modellib = importlib.import_module(model_filename)
		model = None
		target_model_name = model_name.replace('_', '') + 'model'
		for name, cls in modellib.__dict__.items():
		if name.lower() == target_model_name.lower() \
		and issubclass(cls, BaseModel):
		model = cls

		if model is None:
		print("In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase." % (model_filename, target_model_name))
		exit(0)

		return model


		def get_option_setter(model_name):
		"""Return the static method <modify_commandline_options> of the model class."""
		model_class = find_model_using_name(model_name)
		return model_class.modify_commandline_options


		def create_model(opt):
		"""Create a model given the option.

		This function warps the class CustomDatasetDataLoader.
		This is the main interface between this package and 'train.py'/'test.py'

		Example:
		#>>> from models import create_model
		#>>> model = create_model(opt)
		"""
		model = find_model_using_name(opt.model)
		instance = model(opt)
		print("model [%s] was created" % type(instance).__name__)
		return instance

OSAP_web/OSAP_flask/models/Adaptation_model/arcnet_model.py

deleted100644 → 0

+0 −209

Original line number	Diff line number	Diff line
		# -- coding: UTF-8 --
		"""
		@Function:
		@File: arcnet_model.py
		@Date: 2021/11/9 15:20
		@Author: Hever
		"""
		import torch
		import itertools
		from .base_model import BaseModel
		from . import networks
		from models.guided_filter_pytorch.HFC_filter import HFCFilter

		def mul_mask(image, mask):
		return (image + 1) * mask - 1

		class ArcNetModel(BaseModel):
		@staticmethod
		def modify_commandline_options(parser, is_train=True):
		# changing the default values to match the pix2pix paper (https://phillipi.github.io/pix2pix/)
		parser.set_defaults(norm='batch', netG='unet_256', dataset_mode='aligned')
		if is_train:
		parser.set_defaults(pool_size=0, gan_mode='vanilla')
		parser.add_argument('--lambda_L1', type=float, default=100.0, help='weight for L1 loss')
		parser.add_argument('--lambda_L1G', type=float, default=50.0, help='weight for L1G loss')
		parser.add_argument('--lambda_DDP', type=float, default=1, help='weight for DDP')
		parser.add_argument('--lambda_DP', type=float, default=1, help='weight for G loss')

		parser.add_argument('--RMS', action='store_true',)
		parser.add_argument('--filter_width', type=int, default=53, help='weight for G loss')
		parser.add_argument('--nsig', type=int, default=9, help='weight for G loss')
		parser.add_argument('--sub_low_ratio', type=float, default=1.0, help='weight for L1L loss')

		return parser

		def __init__(self, opt):
		"""Initialize the pix2pix class.

		Parameters:
		opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
		"""
		BaseModel.__init__(self, opt)
		self.input_nc = opt.input_nc

		self.loss_names = ['DP', 'DP_fake', 'DP_real',
		'DDP', 'DDP_fake_SB', 'DDP_fake_TB',
		'G', 'G_DP', 'G_L1', 'G_L1G', 'G_DDP']

		self.visual_names = ['real_SA', 'real_SAG', 'fake_SB',
		'fake_SBG', 'real_SB', 'real_SBG',
		'real_TA', 'real_TAG',
		'fake_TB', 'fake_TBG']
		# 初始化guide filter和灰度图工具
		self.hfc_filter = HFCFilter(opt.filter_width, nsig=opt.nsig, sub_low_ratio=opt.sub_low_ratio, sub_mask=True, is_clamp=True).to(self.device)

		if self.isTrain:
		self.model_names = ['G', 'DDP', 'DP']
		else: # during test time, only load G
		self.model_names = ['G']
		self.visual_names = ['real_TA', 'real_TAG',
		'fake_TB', 'fake_TBG']
		# 网络的输出是3个channel
		self.netG = networks.define_G(6, 3, opt.ngf, opt.netG, opt.norm,
		not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)

		# define a discriminator; conditional GANs need to take both input and output images; Therefore, #channels for D is input_nc + output_nc
		if self.isTrain:
		self.netDP = networks.define_D(3, opt.ndf, opt.netD,
		opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
		self.netDDP = networks.define_D(3, opt.ndf, opt.netD,
		opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
		# loss functions
		self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)
		self.criterionL1 = torch.nn.L1Loss()

		# optimizers
		# 增加使用RMS
		if not self.opt.RMS:
		self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
		self.optimizer_D = torch.optim.Adam(itertools.chain(self.netDP.parameters(),
		self.netDDP.parameters()), lr=opt.lr,
		betas=(opt.beta1, 0.999))
		else:
		self.optimizer_G = torch.optim.RMSprop(self.netG.parameters(), lr=opt.lr, alpha=0.9)
		self.optimizer_D = torch.optim.RMSprop(itertools.chain(self.netDP.parameters(),
		self.netDDP.parameters()), lr=opt.lr,
		alpha=0.9)

		self.optimizers.append(self.optimizer_G)
		self.optimizers.append(self.optimizer_D)

		def set_input(self, input):
		"""
		处理输入
		"""
		if self.isTrain:
		AtoB = self.opt.direction == 'AtoB'
		self.real_SA = input['SA' if AtoB else 'SB'].to(self.device)
		self.real_SB = input['SB' if AtoB else 'SA'].to(self.device)
		self.real_TA = input['TA' if AtoB else 'TB'].to(self.device)
		self.S_mask = input['S_mask'].to(self.device)
		self.T_mask = input['T_mask'].to(self.device)

		self.real_SAG = self.hfc_filter(self.real_SA, self.S_mask)
		self.real_TAG = self.hfc_filter(self.real_TA, self.T_mask)
		self.real_SBG = self.hfc_filter(self.real_SB, self.S_mask)

		self.real_SA6 = torch.cat([self.real_SA, self.real_SAG], dim=1)
		self.real_TA6 = torch.cat([self.real_TA, self.real_TAG], dim=1)

		self.image_paths = input['TA_path']
		else:
		AtoB = self.opt.direction == 'AtoB'
		self.real_TA = input['TA' if AtoB else 'TB'].to(self.device)
		self.T_mask = input['T_mask'].to(self.device)
		self.real_TAG = self.hfc_filter(self.real_TA, self.T_mask)
		self.real_TA6 = torch.cat([self.real_TA, self.real_TAG], dim=1)
		self.image_paths = input['TA_path']

		def forward(self):
		"""Run forward pass; called by both functions <optimize_parameters> and <test>."""
		if self.isTrain:
		self.fake_SB = self.netG(self.real_SA6) # G(SA)
		self.fake_SB = mul_mask(self.fake_SB, self.S_mask)
		self.fake_TB = self.netG(self.real_TA6) # G(TA)
		self.fake_TB = mul_mask(self.fake_TB, self.T_mask)

		self.fake_SBG = self.hfc_filter(self.fake_SB, self.S_mask)
		self.fake_SBG = mul_mask(self.fake_SBG, self.S_mask)
		self.fake_TBG = self.hfc_filter(self.fake_TB, self.T_mask)
		self.fake_TBG = mul_mask(self.fake_TBG, self.T_mask)
		else:
		self.fake_TB = self.netG(self.real_TA6) # G(TA)
		self.fake_TB = mul_mask(self.fake_TB, self.T_mask)
		self.fake_TBG = self.hfc_filter(self.fake_TB, self.T_mask)
		self.fake_TBG = mul_mask(self.fake_TBG, self.T_mask)
		# self.fake_SAB = torch.cat((self.real_SA, self.fake_SB), dim=1)
		# self.real_SAB = torch.cat((self.real_SA, self.real_SB), dim=1)


		def backward_DDP(self):
		"""
		Calculate Domain loss for the discriminator, we want to discriminate S and T
		"""
		# Fake Target, detach
		pred_fake_SB = self.netDDP(self.fake_SB.detach())
		pred_fake_TB = self.netDDP(self.fake_TB.detach())

		self.loss_DDP_fake_SB = self.criterionGAN(pred_fake_SB, True) * self.opt.lambda_DDP
		self.loss_DDP_fake_TB = self.criterionGAN(pred_fake_TB, False) * self.opt.lambda_DDP

		# combine loss and calculate gradients
		self.loss_DDP = (self.loss_DDP_fake_SB + self.loss_DDP_fake_TB) * 0.5
		self.loss_DDP.backward()

		def backward_DP(self):
		"""
		Calculate GAN loss for the discriminator
		"""
		pred_fake_SB = self.netDP(self.fake_SB.detach())
		pred_real_SB = self.netDP(self.real_SB.detach())

		self.loss_DP_fake = self.criterionGAN(pred_fake_SB, False)
		self.loss_DP_real = self.criterionGAN(pred_real_SB, True)

		# combine loss and calculate gradients
		self.loss_DP = (self.loss_DP_fake + self.loss_DP_real) * 0.5
		self.loss_DP.backward()


		def backward_G(self):
		"""
		Calculate GAN and L1 loss for the generator
		Generator should fool the DD and DP
		"""
		# First, G(A) should fake the discriminator
		pred_fake_SB = self.netDP(self.fake_SB)

		self.loss_G_DP = self.opt.lambda_DP * self.criterionGAN(pred_fake_SB, True)

		# Second, G(A) = B
		self.loss_G_L1 = self.criterionL1(self.fake_SB, self.real_SB) * self.opt.lambda_L1
		self.loss_G_L1G = self.criterionL1(self.fake_SBG, self.real_SBG) * self.opt.lambda_L1G

		# Third, G(SA) and G(TA) should fool the domain discriminator
		pred_fake_TB = self.netDDP(self.fake_TB)
		pred_fake_SB = self.netDDP(self.fake_SB)
		self.loss_G_DDP = (self.criterionGAN(pred_fake_TB, True) + self.criterionGAN(pred_fake_SB, False)) * \
		0.5 * self.opt.lambda_DDP

		self.loss_G = self.loss_G_DP + self.loss_G_DDP + self.loss_G_L1 + self.loss_G_L1G
		self.loss_G.backward()

		def optimize_parameters(self):
		self.forward() # compute fake images: G(A)
		# update DD (domain discriminator)
		self.set_requires_grad([self.netDP, self.netDDP], True)
		self.optimizer_D.zero_grad() # set D_A and D_B's gradients to zero
		self.backward_DP() # calculate gradients for D_A
		self.backward_DDP() # calculate graidents for D_B
		self.optimizer_D.step() # update D_A and D_B's weights

		# update G
		self.set_requires_grad([self.netDP, self.netDDP],
		False) # Ds require no gradients when optimizing Gs
		self.optimizer_G.zero_grad() # set G_A and G_B's gradients to zero
		self.backward_G() # calculate gradients for G_A and G_B
		self.optimizer_G.step() # update G_A and G_B's weights

OSAP_web/OSAP_flask/models/Adaptation_model/base_model.py

deleted100644 → 0

+0 −232

Original line number	Diff line number	Diff line
		import os
		import torch
		from collections import OrderedDict
		from abc import ABC, abstractmethod
		from . import networks


		class BaseModel(ABC):
		"""This class is an abstract base class (ABC) for models.
		To create a subclass, you need to implement the following five functions:
		-- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).
		-- <set_input>: unpack data from dataset and apply preprocessing.
		-- <forward>: produce intermediate results.
		-- <optimize_parameters>: calculate losses, gradients, and update network weights.
		-- <modify_commandline_options>: (optionally) add model-specific options and set default options.
		"""

		def __init__(self, opt):
		"""Initialize the BaseModel class.

		Parameters:
		opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions

		When creating your custom class, you need to implement your own initialization.
		In this function, you should first call <BaseModel.__init__(self, opt)>
		Then, you need to define four lists:
		-- self.loss_names (str list): specify the training losses that you want to plot and save.
		-- self.model_names (str list): define networks used in our training.
		-- self.visual_names (str list): specify the images that you want to display and save.
		-- self.optimizers (optimizer list): define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example.
		"""
		self.opt = opt
		self.gpu_ids = opt.gpu_ids
		self.isTrain = opt.isTrain
		self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') # get device name: CPU or GPU
		self.save_dir = os.path.join(opt.checkpoints_dir, opt.name) # save all the checkpoints to save_dir
		if opt.preprocess != 'scale_width': # with [scale_width], input images might have different sizes, which hurts the performance of cudnn.benchmark.
		torch.backends.cudnn.benchmark = True
		self.loss_names = []
		self.model_names = []
		self.visual_names = []
		self.optimizers = []
		self.image_paths = []
		self.metric = 0 # used for learning rate policy 'plateau'

		@staticmethod
		def modify_commandline_options(parser, is_train):
		"""Add new model-specific options, and rewrite default values for existing options.

		Parameters:
		parser -- original option parser
		is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.

		Returns:
		the modified parser.
		"""
		return parser

		@abstractmethod
		def set_input(self, input):
		"""Unpack input data from the dataloader and perform necessary pre-processing steps.

		Parameters:
		input (dict): includes the data itself and its metadata information.
		"""
		pass

		@abstractmethod
		def forward(self):
		"""Run forward pass; called by both functions <optimize_parameters> and <test>."""
		pass

		@abstractmethod
		def optimize_parameters(self):
		"""Calculate losses, gradients, and update network weights; called in every training iteration"""
		pass

		def setup(self, opt):
		"""Load and print networks; create schedulers

		Parameters:
		opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
		"""
		if self.isTrain:
		self.schedulers = [networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers]
		if not self.isTrain or opt.continue_train:
		# load_suffix = 'iter_%d' % opt.load_iter if opt.load_iter > 0 else opt.epoch
		load_suffix = '%d' % opt.load_iter if opt.load_iter > 0 else opt.epoch
		self.load_networks(load_suffix)
		self.print_networks(opt.verbose)

		def eval(self):
		"""Make models eval mode during test time"""
		for name in self.model_names:
		if isinstance(name, str):
		net = getattr(self, 'net' + name)
		net.eval()

		def test(self):
		"""Forward function used in test time.

		This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
		It also calls <compute_visuals> to produce additional visualization results
		"""
		with torch.no_grad():
		self.forward()
		self.compute_visuals()

		def compute_visuals(self):
		"""Calculate additional output images for visdom and HTML visualization"""
		pass

		def get_image_paths(self):
		""" Return image paths that are used to load current data"""
		return self.image_paths

		# 使用scheduler规划学习率，退化学习
		def update_learning_rate(self):
		"""Update learning rates for all the networks; called at the end of every epoch"""
		old_lr = self.optimizers[0].param_groups[0]['lr']
		for scheduler in self.schedulers:
		if self.opt.lr_policy == 'plateau':
		scheduler.step(self.metric)
		else:
		scheduler.step()

		lr = self.optimizers[0].param_groups[0]['lr']
		print('learning rate %.7f -> %.7f' % (old_lr, lr))

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

		def get_current_losses(self):
		"""Return traning losses / errors. train.py will print out these errors on console, and save them to a file"""
		errors_ret = OrderedDict()
		for name in self.loss_names:
		if isinstance(name, str):
		errors_ret[name] = float(getattr(self, 'loss_' + name)) # float(...) works for both scalar tensor and float number
		return errors_ret

		def save_networks(self, epoch):
		"""Save all the networks to the disk.

		Parameters:
		epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
		"""
		for name in self.model_names:
		if isinstance(name, str):
		save_filename = '%s_net_%s.pth' % (epoch, name)
		save_path = os.path.join(self.save_dir, save_filename)
		net = getattr(self, 'net' + name)

		if len(self.gpu_ids) > 0 and torch.cuda.is_available():
		torch.save(net.module.cpu().state_dict(), save_path)
		net.cuda(self.gpu_ids[0])
		else:
		torch.save(net.cpu().state_dict(), save_path)

		def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):
		"""Fix InstanceNorm checkpoints incompatibility (prior to 0.4)"""
		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)

		def load_networks(self, epoch):
		"""Load all the networks from the disk.

		Parameters:
		epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
		"""
		for name in self.model_names:
		if isinstance(name, str):
		load_filename = '%s_net_%s.pth' % (epoch, name)
		load_path = os.path.join(self.save_dir, load_filename)
		net = getattr(self, 'net' + name)
		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, map_location=str(self.device))
		if hasattr(state_dict, '_metadata'):
		del state_dict._metadata

		# patch InstanceNorm checkpoints prior to 0.4
		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)

		def print_networks(self, verbose):
		"""Print the total number of parameters in the network and (if verbose) network architecture

		Parameters:
		verbose (bool) -- if verbose: print the network architecture
		"""
		print('---------- Networks initialized -------------')
		for name in self.model_names:
		if isinstance(name, str):
		net = getattr(self, 'net' + 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('-----------------------------------------------')

		def set_requires_grad(self, nets, requires_grad=False):
		"""Set requies_grad=Fasle for all the networks to avoid unnecessary computations
		Parameters:
		nets (network list) -- a list of networks
		requires_grad (bool) -- whether the networks require gradients or not
		"""
		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

OSAP_web/OSAP_flask/models/Adaptation_model/cycle_gan_model.py

deleted100644 → 0

+0 −200

File deleted.

Preview size limit exceeded, changes collapsed.

OSAP_web/OSAP_flask/models/Adaptation_model/guided_filter_pytorch/HFC_filter.py

deleted100644 → 0

+0 −80

File deleted.

Preview size limit exceeded, changes collapsed.

Admin message