def __init__(self, config):
# Config
self.config = config
self.start = 0 # Unless using pre-trained model
# Create directories if not exist
utils.make_folder(self.config.save_path)
utils.make_folder(self.config.model_weights_path)
utils.make_folder(self.config.sample_images_path)
# Copy files
utils.write_config_to_file(self.config, self.config.save_path)
utils.copy_scripts(self.config.save_path)
# Check for CUDA
utils.check_for_CUDA(self)
# Make dataloader
self.dataloader, self.num_of_classes = utils.make_dataloader(
self.config.batch_size_in_gpu, self.config.dataset,
self.config.data_path, self.config.shuffle, self.config.drop_last,
self.config.dataloader_args, self.config.resize,
self.config.imsize, self.config.centercrop,
self.config.centercrop_size)
# Data iterator
self.data_iter = iter(self.dataloader)
# Build G and D
self.build_models()
if self.config.adv_loss == 'dcgan':
self.criterion = nn.BCELoss()
import sys
import utils
from parameters import *
from sagan_models import Generator, Discriminator
if __name__ == '__main__':
config = get_parameters()
config.command = 'python ' + ' '.join(sys.argv)
print(config)
utils.check_for_CUDA(config)
# Load pretrained model (if provided)
if config.pretrained_model != '':
utils.load_pretrained_model(config)
else:
assert config.num_of_classes, "Please provide number of classes! Eg. python3 test.py --num_of_classes 10"
config.G = Generator(config.z_dim, config.g_conv_dim,
config.num_of_classes).to(config.device)
config.D = Discriminator(config.d_conv_dim,
config.num_of_classes).to(config.device)
config.G.eval()
config.D.eval()
print(config.G, config.D)
def __init__(self, config):
# Images data path & Output path
self.dataset = config.dataset
self.data_path = config.data_path
self.save_path = os.path.join(config.save_path, config.name)
# Training settings
self.batch_size = config.batch_size
self.total_step = config.total_step
self.d_steps_per_iter = config.d_steps_per_iter
self.g_steps_per_iter = config.g_steps_per_iter
self.d_lr = config.d_lr
self.g_lr = config.g_lr
self.beta1 = config.beta1
self.beta2 = config.beta2
self.inst_noise_sigma = config.inst_noise_sigma
self.inst_noise_sigma_iters = config.inst_noise_sigma_iters
self.start = 0 # Unless using pre-trained model
# Image transforms
self.shuffle = config.shuffle
self.drop_last = config.drop_last
self.resize = config.resize
self.imsize = config.imsize
self.centercrop = config.centercrop
self.centercrop_size = config.centercrop_size
self.tanh_scale = config.tanh_scale
self.normalize = config.normalize
# Step size
self.log_step = config.log_step
self.sample_step = config.sample_step
self.model_save_step = config.model_save_step
self.save_n_images = config.save_n_images
self.max_frames_per_gif = config.max_frames_per_gif
# Pretrained model
self.pretrained_model = config.pretrained_model
# Misc
self.manual_seed = config.manual_seed
self.disable_cuda = config.disable_cuda
self.parallel = config.parallel
self.dataloader_args = config.dataloader_args
# Output paths
self.model_weights_path = os.path.join(self.save_path,
config.model_weights_dir)
self.sample_path = os.path.join(self.save_path, config.sample_dir)
# Model hyper-parameters
self.adv_loss = config.adv_loss
self.z_dim = config.z_dim
self.g_conv_dim = config.g_conv_dim
self.d_conv_dim = config.d_conv_dim
self.lambda_gp = config.lambda_gp
# Model name
self.name = config.name
# Create directories if not exist
utils.make_folder(self.save_path)
utils.make_folder(self.model_weights_path)
utils.make_folder(self.sample_path)
# Copy files
utils.write_config_to_file(config, self.save_path)
utils.copy_scripts(self.save_path)
# Check for CUDA
utils.check_for_CUDA(self)
# Make dataloader
self.dataloader, self.num_of_classes = utils.make_dataloader(
self.batch_size, self.dataset, self.data_path, self.shuffle,
self.drop_last, self.dataloader_args, self.resize, self.imsize,
self.centercrop, self.centercrop_size)
# Data iterator
self.data_iter = iter(self.dataloader)
# Build G and D
self.build_models()
# Start with pretrained model (if it exists)
if self.pretrained_model != '':
utils.load_pretrained_model(self)
if self.adv_loss == 'dcgan':
self.criterion = nn.BCELoss()
from parameters import get_params
if __name__ == '__main__':
# Get all parameters
args = get_params(eval_mode=True)
args.command = 'python ' + ' '.join(sys.argv)
assert args.pth
assert os.path.exists(args.pth)
assert args.data_path
assert os.path.exists(args.data_path)
assert args.out_path
assert os.path.exists(args.out_path)
# CUDA
utils.check_for_CUDA(args)
# Load pth
pth_dir_name = os.path.dirname(args.pth)
print("Loading model", os.path.join(pth_dir_name, 'model.pth'))
model = torch.load(os.path.join(pth_dir_name, 'model.pth'))
print("Loading model state dict", args.pth)
model.load_state_dict(torch.load(args.pth))
model = model.to(args.device)
# Make dataloader with Eval parameters
print("Making dataloader")
args.valid_split = 0
args.centercrop = False
args.shuffle = False
args.drop_last = False
