def main(args):
# if we enabled GPU mode, set the GPU to use
if args.device_id >= 0:
chainer.cuda.get_device(args.device_id).use()
# Load dataset (we will only use the training set)
if args.mnist:
train, test = chainer.datasets.get_mnist(withlabel=False,
scale=2,
ndim=3)
generator = GeneratorMNIST()
else:
train, test = chainer.datasets.get_cifar10(withlabel=False,
scale=2,
ndim=3)
generator = GeneratorCIFAR()
# subtracting 1, after scaling to 2 (done above) will make all pixels in the range [-1,1]
train -= 1.0
num_training_samples = train.shape[0]
# make data iterators
train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
# build optimizers and models
opt_generator = chainer.optimizers.RMSprop(lr=args.learning_rate)
opt_discriminator = chainer.optimizers.RMSprop(lr=args.learning_rate)
opt_generator.setup(generator)
opt_discriminator.setup(Discriminator())
# make a random noise iterator (uniform noise between -1 and 1)
noise_iter = RandomNoiseIterator(UniformNoiseGenerator(-1, 1, args.num_z),
args.batchsize)
# send to GPU
if args.device_id >= 0:
opt_generator.target.to_gpu()
opt_discriminator.target.to_gpu()
# make the output folder
if not os.path.exists(args.output):
os.makedirs(args.output, exist_ok=True)
print("Starting training loop...")
while train_iter.epoch < args.num_epochs:
training_step(args, train_iter, noise_iter, opt_generator,
opt_discriminator)
print("Finished training.")
def build_model(self):
"""Create a generator and a discriminator."""
# self.data.vertexes
self.G = Generator(self.g_conv_dim, self.z_dim, # 16
31,
self.b_dim, # edges type
self.m_dim, # 4-d vector
self.dropout)
self.D = Discriminator(self.d_conv_dim, self.m_dim, self.b_dim, self.dropout)
self.V = Discriminator(self.d_conv_dim, self.m_dim, self.b_dim, self.dropout)
self.g_optimizer = torch.optim.Adam(list(self.G.parameters())+list(self.V.parameters()),
self.g_lr, [self.beta1, self.beta2])
self.d_optimizer = torch.optim.Adam(self.D.parameters(), self.d_lr, [self.beta1, self.beta2])
if printModel:
self.print_network(self.G, 'G')
self.print_network(self.D, 'D')
self.G.to(self.device)
self.D.to(self.device)
self.V.to(self.device)
def __init__(self, hparams):
super(GAN, self).__init__()
self.hparams = hparams
# networks
self.generator = Generator(self.hparams.latent_dim,
self.hparams.model_dim, 1)
self.discriminator = Discriminator(self.hparams.model_dim,
1,
sigm=True)
self.critic_counter = 0
def build_model(self):
"""Create a generator and a discriminator."""
self.G = Generator(self.g_conv_dim, self.z_dim, self.data.vertexes,
self.data.bond_num_types, self.data.atom_num_types,
self.dropout)
self.D = Discriminator(self.d_conv_dim, self.m_dim, self.b_dim,
self.dropout)
self.V = Discriminator(self.d_conv_dim, self.m_dim, self.b_dim,
self.dropout)
self.g_optimizer = torch.optim.Adam(
list(self.G.parameters()) + list(self.V.parameters()), self.g_lr,
[self.beta1, self.beta2])
self.d_optimizer = torch.optim.Adam(self.D.parameters(), self.d_lr,
[self.beta1, self.beta2])
self.print_network(self.G, 'G')
self.print_network(self.D, 'D')
self.G.to(self.device)
self.D.to(self.device)
self.V.to(self.device)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--gan-dir', type=str, default=None)
parser.add_argument('--ae-dir', type=str, default=None)
parser.add_argument('--start-font',
type=str,
default='LibreBaskerville-Regular')
parser.add_argument('--end-font', type=str, default='Asap-Bold')
args = parser.parse_args()
# Initialize models
gan = {
"gen": Generator().to(device).eval(),
"dis": Discriminator().to(device).eval(),
}
print(f"Loading GAN from {Path(args.gan_dir)}")
gan["gen"].load_state_dict(
torch.load(str(Path(args.gan_dir) / 'gen.ckpt'),
map_location=torch.device('cpu')))
gan["dis"].load_state_dict(
torch.load(str(Path(args.gan_dir) / 'dis.ckpt'),
map_location=torch.device('cpu')))
autoencoder = {
"encoder_A":
ResnetGenerator_3d_conv(input_nc=52, output_nc=52).to(device).eval(),
"encoder_B":
FontEncoder(input_nc=52, output_nc=52).to(device).eval(),
"gen":
FontAdjuster(zdim=32 * 8 * 8).to(device).eval(),
}
print(f"Loading autoencoder from {Path(args.ae_dir)}")
autoencoder["gen"].load_state_dict(
torch.load(str(Path(args.ae_dir) / 'gen.ckpt'),
map_location=torch.device('cpu')))
autoencoder["encoder_A"].load_state_dict(
torch.load(str(Path(args.ae_dir) / 'enc_A.ckpt'),
map_location=torch.device('cpu')))
autoencoder["encoder_B"].load_state_dict(
torch.load(str(Path(args.ae_dir) / 'enc_B.ckpt'),
map_location=torch.device('cpu')))
fonts = [
args.start_font,
args.end_font,
]
dataset = FontDataset('data/jpg', fonts=fonts)
loader = DataLoader(dataset,
batch_size=len(fonts),
sampler=SequentialSampler(dataset))
pipeline(gan, autoencoder, loader)
def build_network(self):
print('[info] Build the network architecture')
self.encoder = Encoder(z_dim=self.opt.latent_dim)
if self.opt.dataset == 'SMPL':
num_verts = 6890
elif self.opt.dataset == 'all_animals':
num_verts = 3889
self.decoder = Decoder(num_verts=num_verts, z_dim=self.opt.latent_dim)
self.discriminator = Discriminator(input_dim=self.opt.latent_dim)
self.encoder.cuda()
self.decoder.cuda()
self.discriminator.cuda()
def _init_models(self):
# Init Model
self.generator = Generator()
self.discriminator = Discriminator(self.config.conv_dim,
self.config.layer_num)
# Init Weights
self.generator.apply(weights_init_normal)
self.discriminator.apply(weights_init_normal)
# Move model to device (GPU or CPU)
self.generator = torch.nn.DataParallel(self.generator).to(self.device)
self.discriminator = torch.nn.DataParallel(self.discriminator).to(
self.device)
def CycleGANmapper(inC, outC, options, init_weights=True):
GA2B = Generator(inC, outC)
DB = Discriminator(outC)
GB2A = Generator(outC, inC)
DA = Discriminator(inC)
if options["cuda"]:
GA2B.cuda()
DB.cuda()
GB2A.cuda()
DA.cuda()
if init_weights:
GA2B.apply(weights_init_normal)
DB.apply(weights_init_normal)
GB2A.apply(weights_init_normal)
DA.apply(weights_init_normal)
return (GA2B, DB), (GB2A, DA)
if options["continued"]:
GA2B.load_state_dict(torch.load("output/GEN_AtoB.pth"))
DB.load_state_dict(torch.load("output/DIS_B.pth"))
GB2A.load_state_dict(torch.load("output/GEN_BtoA.pth"))
DA.load_state_dict(torch.load("output/DIS_A.pth"))
return (GA2B, DB), (GB2A, DA)
def main(args):
args.latent_size = 64
args.image_size = 784
args.hidden_size = 256
""" Create a directory if not exists """
if not os.path.isdir('checkpoints'):
os.mkdir('checkpoints')
if not os.path.isdir(args.data_path):
os.mkdir(args.data_path)
if not os.path.isdir(args.save_path):
os.mkdir(args.save_path)
""" Settings for training """
torch.manual_seed(args.seed)
""" Device configuration """
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
multi_gpu = torch.cuda.device_count() > 1
print("[Info] The system consist of {} {}.".format(torch.cuda.device_count() if device == 'cuda' else '1', device))
""" Prepare the dataloaders """
train_loader, valid_loader, train_size, valid_size = prepare_dataloaders(args)
""" Define the generator & discriminator. """
generator = Generator(args).to(device)
discriminator = Discriminator(args).to(device)
models = [generator, discriminator]
""" Define the loss function & optimizers """
criterion = torch.nn.BCELoss().to(device)
g_optimizer = torch.optim.Adam(generator.parameters(), lr=args.lr)
d_optimizer = torch.optim.Adam(discriminator.parameters(), lr=args.lr)
optimizers = [g_optimizer, d_optimizer]
""" Define the learning rate scheduler """
# scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
""" Iterate over the data. """
best_loss = float('inf')
for epoch in range(1, args.epochs + 1):
# adjust_learning_rate(optimizers, epoch)
# train for one epochs
_ = run_epoch(models, 'train', train_loader, train_size, args, epoch, criterion, optimizers, device)
with torch.no_grad():
loss = run_epoch(models, 'valid', valid_loader, valid_size, args, epoch, criterion, optimizers, device)
if best_loss > loss:
best_loss = loss
print("[Info] Save the model. : {}".format(epoch))
torch.save(models[0].state_dict(), './checkpoints/G.ckpt')
torch.save(models[1].state_dict(), './checkpoints/D.ckpt')
def main():
config = Config()
train_iters, dev_iters, test_iters, vocab = load_dataset(config,
train_pos='cvet2.train', train_neg='push2.train',
dev_pos='cvet2.dev', dev_neg='push2.dev',
test_pos='cvet2.test', test_neg='push2.test')
print('Vocab size:', len(vocab))
model_F = StyleTransformer(config, vocab).to(config.device)
model_D = Discriminator(config, vocab).to(config.device)
print(config.discriminator_method)
train(config, vocab, model_F, model_D, train_iters, dev_iters, test_iters)
torch.save(model_F.state_dict(), 'modelF_trained')
torch.save(model_D.state_dict(), 'modelD_trained')
def network_initializers(self, hr_shape, use_LeakyReLU_Mish=False):
generator = GeneratorRRDB(self.opt.channels,
filters=64,
num_res_blocks=self.opt.residual_blocks,
use_LeakyReLU_Mish=use_LeakyReLU_Mish).to(
self.device, non_blocking=True)
discriminator = Discriminator(
input_shape=(self.opt.channels, *hr_shape),
use_LeakyReLU_Mish=use_LeakyReLU_Mish).to(self.device,
non_blocking=True)
feature_extractor = FeatureExtractor().to(self.device,
non_blocking=True)
# Set feature extractor to inference mode
feature_extractor.eval()
return discriminator, feature_extractor, generator
def main():
config = Config()
train_iters, test_iters, vocab = load_enron(config)
print('Vocab size:', len(vocab))
model_F = StyleTransformer(config, vocab).to(config.device)
model_D = Discriminator(config, vocab).to(config.device)
print(config.discriminator_method)
# last_checkpoint = most_recent_path(most_recent_path(config.save_path), return_two=True)
# if last_checkpoint:
# print(last_checkpoint)
# model_D.load_state_dict(torch.load(last_checkpoint[1]))
# model_F.load_state_dict(torch.load(last_checkpoint[0]))
train(config, vocab, model_F, model_D, train_iters, test_iters)
def __init__(self):
super(DCGAN, self).__init__()
# Define Network
self.netG = Generator(_C)
self.netD = Discriminator(_C)
## Define labels
self.real_label = 1
self.fake_label = 0
self.gen_img = None
## Initialize Weight
self.netD.apply(weights_init)
self.netG.apply(weights_init)
def __init__(self):
self.generator = Generator()
self.discriminator = Discriminator()
self.loss = nn.MSELoss()
self.generator.cuda()
self.discriminator.cuda()
self.loss.cuda()
self.optimizer_g = optim.Adam(self.generator.parameters(),
lr=config.lr,
betas=(config.beta, 0.999))
self.optimizer_d = optim.Adam(self.discriminator.parameters(),
lr=config.lr,
betas=(config.beta, 0.999))
def _set_model(self, device, hr_shape):
# Initialize generator and discriminator
self.generator = GeneratorRRDB(
opt.channels, filters=64,
num_res_blocks=opt.residual_blocks).to(device)
self.discriminator = Discriminator(input_shape=(opt.channels,
*hr_shape)).to(device)
self.feature_extractor = FeatureExtractor().to(device)
# Set feature extractor to inference mode
self.feature_extractor.eval()
# Losses
self.criterion_GAN = torch.nn.BCEWithLogitsLoss().to(device)
self.criterion_content = torch.nn.L1Loss().to(device)
self.criterion_pixel = torch.nn.L1Loss().to(device)
def office():
init_random_seed(params.manual_seed)
# load dataset
src_data_loader = get_data_loader(params.src_dataset)
src_data_loader_eval = get_data_loader(params.src_dataset, train=False)
tgt_data_loader = get_data_loader(params.tgt_dataset)
tgt_data_loader_eval = get_data_loader(params.tgt_dataset, train=False)
# load models
src_encoder = init_model(net=LeNetEncoder(),
restore=params.src_encoder_restore)
src_classifier = init_model(net=LeNetClassifier(),
restore=params.src_classifier_restore)
tgt_encoder = init_model(net=LeNetEncoder(),
restore=params.tgt_encoder_restore)
critic = init_model(Discriminator(input_dims=params.d_input_dims,
hidden_dims=params.d_hidden_dims,
output_dims=params.d_output_dims),
restore=params.d_model_restore)
if not (src_encoder.restored and src_classifier.restored and
params.src_model_trained):
src_encoder, src_classifier = train_src(
src_encoder, src_classifier, src_data_loader)
# eval source model
# print("=== Evaluating classifier for source domain ===")
# eval_src(src_encoder, src_classifier, src_data_loader_eval)
# train target encoder by GAN
# init weights of target encoder with those of source encoder
if not tgt_encoder.restored:
tgt_encoder.load_state_dict(src_encoder.state_dict())
if not (tgt_encoder.restored and critic.restored and
params.tgt_model_trained):
tgt_encoder = train_tgt(src_encoder, tgt_encoder, critic,
src_data_loader, tgt_data_loader)
# eval target encoder on test set of target dataset
print(">>> domain adaption <<<")
acc = eval_tgt(tgt_encoder, src_classifier, tgt_data_loader_eval)
return acc
def group(conv):
d_model_restore = "snapshots//CONV_" + str(
conv) + "_ACTIVATIONS-ADDA-critic-final.pt"
src_classifier_restore = "snapshots//CONV_" + str(
conv) + "_ACTIVATIONS-ADDA-source-classifier-final.pt"
tgt_encoder_restore = "snapshots//CONV_" + str(
conv) + "_ACTIVATIONS-ADDA-target-encoder-final.pt"
src_encoder_restore = "snapshots//CONV_" + str(
conv) + "_ACTIVATIONS-ADDA-source-encoder-final.pt"
if int(conv) == 1:
src_encoder = init_model(net=AurielEncoder(),
restore=src_encoder_restore)
src_classifier = init_model(net=AurielClassifier(),
restore=src_classifier_restore)
tgt_encoder = init_model(net=AurielEncoder(),
restore=tgt_encoder_restore)
elif int(conv) == 2:
src_encoder = init_model(net=BeatriceEncoder(),
restore=src_encoder_restore)
src_classifier = init_model(net=BeatriceClassifier(),
restore=src_classifier_restore)
tgt_encoder = init_model(net=BeatriceEncoder(),
restore=tgt_encoder_restore)
elif int(conv) == 3:
src_encoder = init_model(net=CielEncoder(),
restore=src_encoder_restore)
src_classifier = init_model(net=CielClassifier(),
restore=src_classifier_restore)
tgt_encoder = init_model(net=CielEncoder(),
restore=tgt_encoder_restore)
elif int(conv) == 4:
src_encoder = init_model(net=DioneEncoder(),
restore=src_encoder_restore)
src_classifier = init_model(net=DioneClassifier(),
restore=src_classifier_restore)
tgt_encoder = init_model(net=DioneEncoder(),
restore=tgt_encoder_restore)
else:
raise RuntimeError("conv number must be 1, 2, 3, or 4.")
critic = init_model(Discriminator(input_dims=params.d_input_dims,
hidden_dims=params.d_hidden_dims,
output_dims=params.d_output_dims),
restore=d_model_restore)
return src_encoder, src_classifier, tgt_encoder, critic
def main():
config = Config()
parser = argparse.ArgumentParser()
add_generic_args(parser, os.getcwd())
parser = BartSystem.add_model_specific_args(parser, os.getcwd())
args = parser.parse_args()
# Some values from Config class needs to be copied to args to work.
setattr(config, "num_train_epochs", args.num_train_epochs)
setattr(config, "save_path", args.output_dir)
setattr(args, "learning_rate", config.lr_F)
# Create output directory.
timestamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
setattr(config, "save_folder", os.path.join(config.save_path, timestamp))
os.makedirs(os.path.join(config.save_folder, 'ckpts'))
init_logger(config.save_folder)
logger = logging.getLogger(__name__)
model_F = BartSystem(args).to(config.device)
# Don't use the trainer to fit the model
args.do_train = False
# trainer = generic_train(model_F, args)
if args.output_dir:
try:
checkpoints = list(
sorted(
glob.glob(os.path.join(args.output_dir,
"checkpointepoch=*.ckpt"),
recursive=True)))
if checkpoints[-1]:
BartSystem.load_from_checkpoint(checkpoints[-1])
logger.info("Load checkpoint sucessfully!")
except:
logger.info("Failed to load checkpoint!")
# train_iters, dev_iters, test_iters, vocab = load_dataset(config)
train_iters, dev_iters, test_iters = model_F.train_dataloader(
), model_F.val_dataloader(), model_F.test_dataloader()
model_D = Discriminator(config, model_F.tokenizer).to(config.device)
logger.info(config.discriminator_method)
# import pdb
# pdb.set_trace()
logger.info(model_D)
train(config, model_F, model_D, train_iters, dev_iters, test_iters)
def trio(tgt_classifier_net, tgt_encoder_net, src_dataset, tgt_dataset, conv):
print('loading pretrained trio after conv ' + str(conv) + '...')
tgt_classifier = init_model(net=tgt_classifier_net,
restore= str(conv) + "_snapshots/" + \
src_dataset + "-ADDA-target-classifier-final.pt")
tgt_encoder = init_model(net=tgt_encoder_net,
restore= str(conv) + "_snapshots/" + \
tgt_dataset + "-ADDA-target-classifier-final.pt")
critic = init_model(Discriminator(input_dims=params.d_input_dims,
hidden_dims=params.d_hidden_dims,
output_dims=params.d_output_dims),
restore= str(conv) + "_snapshots/" + \
tgt_dataset + "-ADDA-target-classifier-final.pt")
return tgt_classifier, tgt_encoder, critic
def main():
torch.random.manual_seed(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
parser = argparse.ArgumentParser()
parser.add_argument('--root', type=str, default='data')
parser.add_argument('--batch-size', type=int, default=128)
parser.add_argument('--lr', type=float, default=2e-4)
parser.add_argument('--workers', type=int, default=0)
parser.add_argument('--epochs', type=int, default=20)
config = parser.parse_args()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# networks
net_g = Generator().to(device)
net_d = Discriminator().to(device)
print(net_g)
print(net_d)
# optimizer
optimizer_g = optim.Adam(net_g.parameters(),
lr=config.lr,
betas=(0.5, 0.999))
optimizer_d = optim.Adam(net_d.parameters(),
lr=config.lr,
betas=(0.5, 0.999))
print(optimizer_d)
print(optimizer_g)
# data loader
dataloader = get_loader(config.root, config.batch_size, config.workers)
trainer = Trainer(net_g, net_d, optimizer_g, optimizer_d, dataloader,
device)
plotter = PlotHelper('samples/loss.html')
for epoch in range(config.epochs):
loss_g, loss_d = trainer.train()
print('Train epoch: {}/{},'.format(epoch + 1, config.epochs),
'loss g: {:.6f}, loss d: {:.6f}.'.format(loss_g, loss_d))
trainer.save_sample('samples/sample_{:02d}.jpg'.format(epoch + 1))
plotter.append(loss_g, loss_d, epoch + 1)
def main(config):
logger = prepare_logger(config)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# get loaders
if not config.is_train_source:
target_loader = get_loader(type="MNIST",
train=False,
batch_size=config.batch_size)
source_train_loader = get_loader(type="SVHN",
train=True,
batch_size=config.batch_size)
source_test_loader = get_loader(type="SVHN",
train=False,
batch_size=config.batch_size)
# build source classifier
model_src = LeNet(config.num_gpus).to(device)
if (not config.is_train_source) or config.is_finetune:
model_src.load_state_dict(torch.load(config.model_dir))
# train source classifier
if config.is_train_source:
logger.info("train source classifier..")
train_source(model_src, source_train_loader, source_test_loader,
config, logger)
logger.info("evaluate source classifier..")
logger.info("test accurracy in source domain: %f\n" %
(evaluate(model_src, source_test_loader)))
else:
# initialize target classifer with source classifer
model_trg = torch.load(open("./pretrained/lenet-source.pth", "rb"))
# build discriminator
D = Discriminator(config.num_gpus)
# adaptation process
logger.info("start adaptation process..")
adapt_target_domain(D, model_src, model_trg, source_train_loader,
target_loader, config)
logger.info("evaluate target classifier..")
logger.info("accurracy in target domain: %f\n" %
(evaluate(model_trg, target_loader)))
def load_discriminator(data=None,
discriminator_weights=None,
backbone_weights='backbone_posttrained_weights.h5',
clear_session=True):
if (clear_session):
keras.backend.clear_session()
backbone = ResNet()
backbone(data.get_test()[0])
discriminator = Discriminator(backbone)
if (discriminator_weights):
discriminator.load_weights(discriminator_weights)
if (backbone_weights):
backbone.load_weights(backbone_weights)
discriminator(data.get_test()[0])
return discriminator
def __init__(self, cfg):
super(SRGAN, self).__init__()
# Training stats
self.global_step = tf.train.get_or_create_global_step()
self.epoch = tf.Variable(0)
# Load models
self.generator = Generator(cfg)
self.discriminator = Discriminator(cfg)
#self.vgg = tf.keras.applications.vgg19.VGG19(include_top=False, weights='imagenet', input_shape=cfg.hr_resolution + (3,), pooling=None)
self.gen_optim = tf.train.AdamOptimizer(cfg.learning_rate)
self.disc_optim = tf.train.AdamOptimizer(cfg.learning_rate)
self.build_writers()
self.preprocessing()
self.pretrain()
def _init_models(self, mode: str):
""" Initialises models, loss functions, optimisers and sets models to training mode """
# Task Learner
self.task_learner = TaskLearner(**self.model_config['TaskLearner']['Parameters'],
vocab_size=self.vocab_size,
tagset_size=self.tagset_size,
task_type=self.task_type).to(self.device)
# Loss functions
if self.task_type == 'SEQ':
self.tl_loss_fn = nn.NLLLoss().to(self.device)
if self.task_type == 'CLF':
self.tl_loss_fn = nn.CrossEntropyLoss().to(self.device)
# Optimisers
self.tl_optim = optim.SGD(self.task_learner.parameters(),
lr=self.model_config['TaskLearner']['learning_rate'])#, momentum=0, weight_decay=0.1)
# Learning rate scheduler
# Note: LR likely GT Adam
# self.tl_sched = optim.lr_scheduler.ReduceLROnPlateau(self.tl_optim, 'min', factor=0.5, patience=10)
# Training Modes
self.task_learner.train()
# SVAAL needs to initialise SVAE and DISC in addition to TL
if mode == 'svaal':
# Models
self.svae = SVAE(**self.model_config['SVAE']['Parameters'],
vocab_size=self.vocab_size).to(self.device)
self.discriminator = Discriminator(**self.model_config['Discriminator']['Parameters']).to(self.device)
# Loss Function (SVAE defined within its class)
self.dsc_loss_fn = nn.BCELoss().to(self.device)
# Optimisers
# Note: Adam will likely have a lower lr than SGD
self.svae_optim = optim.Adam(self.svae.parameters(),
lr=self.model_config['SVAE']['learning_rate'])
self.dsc_optim = optim.Adam(self.discriminator.parameters(),
lr=self.model_config['Discriminator']['learning_rate'])
# Training Modes
self.svae.train()
self.discriminator.train()
print(f'{datetime.now()}: Models initialised successfully')
def __init__(self, _argv):
self.init_parser(_argv)
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# init dataloader
trans = transforms.RandomHorizontalFlip(p=0.5)
self.ds = ArtDataset(self.parser.data_dir, transform=trans)
self.data_loader = DataLoader(self.ds, batch_size=self.parser.batch_size, shuffle=True, num_workers=8, drop_last=True)
# self.data is an infinite iterator which will return the next batch and loop back to the start
self.data = self.get_infinite_batches(self.data_loader)
# Load and initialize the generator and discriminator models
self.G = Generator(self.parser.lat_dim, self.parser.n_classes, self.parser.cont_dim, self.parser.img_sz)
self.D = Discriminator(self.parser.lat_dim, self.parser.n_classes, self.parser.cont_dim, self.parser.img_sz)
self.G.train()
self.D.train()
self.G.to(self.device)
self.D.to(self.device)
self.init_weights(self.G)
self.init_weights(self.D)
# optimizers
#value with wgan paper
self.weight_cliping_limit = 0.02
self.G_opt = optim.Adam(self.G.parameters(), lr=self.parser.G_lr, betas=(0.5, 0.9))
self.D_opt = optim.Adam(self.D.parameters(), lr=self.parser.D_lr, betas=(0.5, 0.9))
self.critic_iter = 10
self.g_start_batch_num = 0
# self.G_scheduler = lr_scheduler.StepLR(self.G_opt, step_size=10, gamma=0.1, last_epoch=-1)
# self.D_scheduler = lr_scheduler.StepLR(self.D_opt, step_size=10, gamma=0.1, last_epoch=-1)
# loss functions and weights
self.adversarial_loss = torch.nn.BCELoss()
self.categorical_loss = torch.nn.CrossEntropyLoss()
self.continuous_loss = torch.nn.MSELoss()
self.lambda_cat = 1
self.lambda_cont = 0.1
#static codes for sampling
# Static sample
self.static_noise = Variable(torch.randn(5**2, self.parser.lat_dim)).to(self.device)
self.static_cont_code = Variable(torch.rand(5**2, self.parser.cont_dim)).to(self.device)
self.static_class_code = self.gen_class_code(5**2, self.parser.n_classes).to(self.device)
def models(channels):
"""
Creates and initializes the models
:return: Encoder, Generator, Discriminator
"""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
encoder = Encoder(channels).to(device)
encoder.apply(init_weights)
generator = Generator(channels).to(device)
generator.apply(init_weights)
discriminator = Discriminator(channels).to(device)
discriminator.apply(init_weights)
return encoder, generator, discriminator
def __init__(self, device, trainData, validData, hidden_size, lr,
batch_size, arch):
self.device = device
self.trainData = trainData
self.validData = validData
self.model = SimpleNet(hidden_size).to(device)
self.generator = Generator(hidden_size).to(device)
self.discriminator = Discriminator(hidden_size).to(device)
self.opt = torch.optim.Adam(self.model.parameters(), lr=1e-3)
self.opt_G = torch.optim.Adam(self.generator.parameters(), lr=1e-4)
self.opt_D = torch.optim.Adam(self.discriminator.parameters(), lr=1e-4)
self.criterion = torch.nn.BCEWithLogitsLoss()
self.scheduler = StepLR(self.opt, step_size=150, gamma=0.5)
self.scheduler_G = StepLR(self.opt_G, step_size=300, gamma=0.5)
self.scheduler_D = StepLR(self.opt_D, step_size=300, gamma=0.5)
self.batch_size = batch_size
self.arch = arch
self.history = {'train': [], 'valid': []}
def init_models(args):
netG = Generator(args).cuda()
netD = Discriminator(args).cuda()
netL = PreActResNet().cuda()
print (netG, netD)
optimG = torch.optim.Adam(netG.parameters(), betas=(0.5, 0.999), lr=args.lr)
optimD = torch.optim.Adam(netD.parameters(), betas=(0.5, 0.999), lr=args.lr)
optimL = torch.optim.Adam(netL.parameters(), betas=(0.5, 0.999), lr=args.lr)
if args.resume:
Gpath = 'experiments/{}/models/netG_{}.pt'.format(args.name, args.load_step)
Dpath = 'experiments/{}/models/netD_{}.pt'.format(args.name, args.load_step)
Lpath = 'experiments/{}/models/netL_{}.pt'.format(args.name, args.load_step)
netG, optimG = utils.load_model(args, netG, optimG, Gpath)
netD, optimD = utils.load_model(args, netD, optimD, Dpath)
netL, optimL = utils.load_model(args, netL, optimL, Lpath)
return (netG, optimG), (netD, optimD), (netL, optimL)
def init_models(args):
netG = Generator(args).cuda()
netD = Discriminator(args).cuda()
print(netG, netD)
optimG = torch.optim.Adam(netG.parameters(),
betas=(0.5, 0.999),
lr=args.lr)
optimD = torch.optim.Adam(netD.parameters(),
betas=(0.5, 0.999),
lr=args.lr)
str = 'experiments/{}/models/'.format(args.name)
if args.resume_G is not None:
netG, optimG = utils.load_model(netG, optimG, str + args.resume_G)
if args.resume_D is not None:
netD, optimD = utils.load_model(netG, optimD, str + args.resume_D)
return (netG, optimG), (netD, optimD)
def test_coupled_weights_of_backbone():
"""
This function will fail because there are multiple models defined
in the keras/tensorflow graph which are not used during training.
Returns:
bool -- [description]
"""
data = mnist_data()
backbone = ResNet()
preds = backbone(data.get_test()[0])
gen = ResGen(backbone)
input_shape = gen.get_input_shape()
rand_data_shape = ((50, ) + input_shape[1:] + (1, ))
random_noise_data = np.random.normal(size=rand_data_shape)
discriminator = Discriminator(backbone)
classifier = Classifier(backbone, 10)
discriminator_predicitons_1 = discriminator(data.get_test()[0])
classifier_predicitons_1 = classifier.predict(data.get_test()[0])
generator_predictions_1 = gen.predict(random_noise_data)[0]
classifier.compile(optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
classifier.summary()
# classifier.fit(x=x_train,y=y_train,batch_size=6000,epochs=1, validation_data=(x_vali,y_vali),callbacks=[checkpoint])
classifier.fit(x=data.get_n_samples(35)[0],
y=data.get_n_samples(35)[1],
batch_size=6000,
epochs=1,
validation_data=data.get_vali())
discriminator_predicitons_2 = discriminator(data.get_test()[0])
classifier_predicitons_2 = classifier.predict(data.get_test()[0])
generator_predictions_2 = gen.predict(random_noise_data)[0]
discriminator_diff = discriminator_predicitons_1 - discriminator_predicitons_2
classifier_diff = classifier_predicitons_1 - classifier_predicitons_2
generator_diff = generator_predicitons_1 - generator_predicitons_2
return True