def build_model_and_trainer(config, data_loader):
if config.model.type == 'classifier':
model_builder = Classifier(config)
model, parallel_model = WithLoadWeights(model_builder, model_name='classifier') \
.build_model(model_name='classifier')
trainer = ClassifierTrainer(model, parallel_model, data_loader, config)
return model, trainer
elif config.model.type == 'dcgan':
g_model_builder = Generator(config)
d_model_builder = Discriminator(config)
# c_model_builder = Classifier(config)
g = g_model_builder.define_model('generator')
d, parallel_d = d_model_builder.build_model('discriminator')
# c, _ = c_model_builder.build_model('classifier')
# Load weights to classifier
# checkpoint_path = './experiments/classifier_mnist/checkpoints/0050-classifier.hdf5'
# if os.path.exists(checkpoint_path):
# c.load_weights(checkpoint_path)
combined_model_builder = GANCombined(config)
combined, parallel_combined = WithLoadWeights(combined_model_builder, model_name='combined') \
.build_model(g=g, d=d, model_name='combined')
# .build_model(g=g, d=d, c=c, model_name='combined')
# trainer = GANTrainer(data_loader, config, g, d, parallel_d, c, combined, parallel_combined)
trainer = GANTrainer(data_loader, config, g, d, parallel_d, combined,
parallel_combined)
return combined, trainer
class GAN(nn.Module):
def __init__(self,
num_fakes=16,
rand_dim=128,
size=(32, 32),
use_gpu=False):
super(GAN, self).__init__()
self.generator = Generator(num_fakes, rand_dim, size, use_gpu)
self.discriminator = Discriminator((1, ) + size)
def forward(self, x=None, mode='D'):
if mode == 'D':
with torch.no_grad():
fake_imgs = self.generator()
self.discriminator.freeze(False)
fake_imgs = fake_imgs.detach()
imgs = torch.cat([x, fake_imgs], dim=0)
out = self.discriminator(imgs)
return out
else:
self.discriminator.freeze(True)
fake_imgs = self.generator()
out = self.discriminator(fake_imgs)
return fake_imgs, out
def init_models(self):
self.G_X = Generator(3, 3, nn.InstanceNorm2d)
self.D_X = Discriminator(3)
self.G_Y = Generator(3, 3, nn.InstanceNorm2d)
self.D_Y = Discriminator(3)
if self.load_models:
self.G_X.load_state_dict(
torch.load(self.output_path + "models/G_X",
map_location='cpu'))
self.G_Y.load_state_dict(
torch.load(self.output_path + "models/G_Y",
map_location='cpu'))
self.D_X.load_state_dict(
torch.load(self.output_path + "models/D_X",
map_location='cpu'))
self.D_Y.load_state_dict(
torch.load(self.output_path + "models/D_Y",
map_location='cpu'))
else:
self.G_X.apply(init_func)
self.G_Y.apply(init_func)
self.D_X.apply(init_func)
self.D_Y.apply(init_func)
self.G_X.to(self.device)
self.G_Y.to(self.device)
self.D_X.to(self.device)
self.D_Y.to(self.device)
def make_model(args, dataset):
N = args.gpu_num
class Model(): pass
model = Model()
netG = Generator(args)
netD = Discriminator(args)
netRNN = DS_SJE(args, dataset.embed_mat)
if N == 0:
setattr(model, "netG_0", netG)
setattr(model, "netD_0", netD)
setattr(model, "netRNN_0", netRNN)
else:
for i in range(N):
temp_netG = netG.copy()
temp_netD = netD.copy()
temp_netRNN = netRNN.copy()
temp_netG.gpu_id = i
temp_netD.gpu_id = i
temp_netRNN.gpu_id = i
setattr(model, f"netG_{i}", temp_netG.to_gpu(i))
setattr(model, f"netD_{i}", temp_netD.to_gpu(i))
setattr(model, f"netRNN_{i}", temp_netRNN.to_gpu(i))
return model
def __init__(self,
num_fakes=16,
rand_dim=128,
size=(32, 32),
use_gpu=False):
super(GAN, self).__init__()
self.generator = Generator(num_fakes, rand_dim, size, use_gpu)
self.discriminator = Discriminator((1, ) + size)
def __init__(self):
self.env = gym.make('CartPole-v1')
self.expert_traj_fpath = os.path.join('GAIL', 'expert_trajectories', 'expert_traj_test.npz')
self.save_policy_fpath = os.path.join('GAIL','gail_actor1.pt')
self.save_rewards_fig = os.path.join('GAIL', 'gail_rewards.png')
#
self.state_space = 4
self.action_space = 2
# These values are taken from the env's state and action sizes
self.actor_net = ActorNet(self.state_space, self.action_space)
self.actor_net.to(device=Device)
self.actor_optim = torch.optim.Adam(self.actor_net.parameters(), lr = 0.0001)
self.critic_net = CriticNet(self.state_space)
self.critic_net.to(Device)
self.critic_net_optim = torch.optim.Adam(self.critic_net.parameters(), lr = 0.0001)
self.discriminator = Discriminator(self.state_space, self.action_space)
self.discriminator.to(Device)
self.discriminator_optim = torch.optim.Adam(self.discriminator.parameters(), lr = 0.0001)
# Storing all the values used to calculate the model losses
self.traj_obs = []
self.traj_actions = []
self.traj_rewards = []
self.traj_dones = []
self.traj_logprobs = []
self.traj_logits = []
self.traj_state_values = []
# Discount factor
self.gamma = 0.95
# Bias Variance tradeoff (higher value results in high variance, low bias)
self.gae_lambda = 0.95
# These two will be used during the training of the policy
self.ppo_batch_size = 500
self.ppo_epochs = 12
self.ppo_eps = 0.2
# Discriminator
self.num_expert_transitions = 150
# These will be used for the agent to play using the current policy
self.max_eps = 5
# Max steps in mountaincar ex is 200
self.max_steps = 800
# documenting the stats
self.avg_over = 5 # episodes
self.stats = {'episode': 0, 'ep_rew': []}
def build_model(self):
self.G = Generator(self.noise_n, self.G_last_act)
self.D = Discriminator(self.D_out_n, self.D_last_act)
self.G_optimizer = torch.optim.Adam(self.G.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 torch.cuda.is_available():
self.G.cuda()
self.D.cuda()
def build_discriminator_trainer(self):
"""
Build a Keras model for training image and mask discriminators.
"""
# Mask Discriminator
D_Mask = Discriminator(self.conf.input_shape,
output='2D',
downsample_blocks=3,
name='D_M')
D_Mask.build()
self.MaskDiscriminator = D_Mask.model
real_M = Input(self.conf.input_shape)
fake_M = Input(self.conf.input_shape)
dis_real_M = self.MaskDiscriminator(real_M)
dis_fake_M = self.MaskDiscriminator(fake_M)
D_Image = Discriminator(self.conf.input_shape,
output='2D',
downsample_blocks=3,
name='D_X')
D_Image.build()
self.ImageDiscriminator = D_Image.model
real_X = Input(self.conf.input_shape)
fake_X = Input(self.conf.input_shape)
dis_real_X = self.ImageDiscriminator(real_X)
dis_fake_X = self.ImageDiscriminator(fake_X)
self.D_model = Model(
inputs=[real_M, fake_M, real_X, fake_X],
outputs=[dis_real_M, dis_fake_M, dis_real_X, dis_fake_X])
self.D_model.compile(Adam(lr=0.0001, beta_1=0.5), loss='mse')
log.info('Discriminators Trainer')
self.D_model.summary(print_fn=log.info)
def get_modules(opt):
modules = {}
disc = Discriminator()
gen = Generator()
clf = Classifier()
if opt.cuda:
disc = disc.cuda()
gen = gen.cuda()
clf = clf.cuda()
modules['Discriminator'] = disc
modules['Generator'] = gen
modules['Classifier'] = clf
return modules
def test_adapt():
args = parse()
print(args)
dataset = SepeDataset(args.poses_train,args.images_train,coor_layer_flag =False)
dataloader = DataLoader(dataset, batch_size=1,shuffle=True ,num_workers=1,drop_last=True,worker_init_fn=lambda wid:np.random.seed(np.uint32(torch.initial_seed() + wid)))
dataset_tgt = SepeDataset(args.poses_target,args.images_target,coor_layer_flag =False)
dataloader_tgt = DataLoader(dataset_tgt, batch_size=1,shuffle=True ,num_workers=1,drop_last=True,worker_init_fn=lambda wid:np.random.seed(np.uint32(torch.initial_seed() + wid)))
src_extractor = DVOFeature()
tgt_extractor = DVOFeature()
src_extractor.load_state_dict(torch.load(args.feature_model))
tgt_extractor.load_state_dict(torch.load(args.feature_model))
dvo_discriminator = Discriminator(500,500,2)
adapt(src_extractor,tgt_extractor,dvo_discriminator,dataloader,dataloader_tgt,args)
torch.save(tgt_extractor.state_dict(),'tgt_feature_'+args.tag+str(args.epoch)+'.pt')
torch.save(dvo_discriminator.state_dict(),'dis_'+args.tag+str(args.epoch)+'.pt')
def get_energy(args, input_dims, device):
if args.discriminator == 'convolutional':
return Discriminator(nn_type=args.d_model,
bn=args.bn,
no_trunc=args.no_trunc,
skipinit=args.skipinit).to(device)
elif args.discriminator == 'nvp':
return energy_model.NVP([input_dims],
device,
args.num_blocks,
mode='discriminator',
with_bn=args.dis_bn).to(device)
elif args.discriminator == 'made':
return energy_model.MADEGenerator([input_dims],
mode='discriminator').to(device)
elif args.discriminator == 'maf':
return energy_model.FlowGenerator([input_dims],
device,
args.num_blocks,
'maf',
mode='discriminator',
with_bn=args.dis_bn).to(device)
elif args.discriminator == 'mogmaf':
return energy_model.FlowGenerator([input_dims],
device,
args.num_blocks,
'mogmaf',
mode='discriminator',
with_bn=args.dis_bn).to(device)
elif args.discriminator == 'toy':
return tm.Discriminator(3)
elif args.discriminator == 'are':
return energy_model.Discriminator(input_dims, device).to(device)
elif args.discriminator == 'are4':
return energy_model.Discriminator4(input_dims, device).to(device)
def __init__(self, args):
self.args = args
# To write to Tensorboard
self.writer = SummaryWriter()
# Holds all data classes needed
self.data = GANData(args, root='./data')
# Instantiate models and load to device
self.D = Discriminator(784, args.d_hidden_size).to(args.device)
self.G = Generator(784, args.latent_size, args.g_hidden_size).to(args.device)
# Instantiate criterion used for both D and G
self.criterion = nn.BCELoss()
# Instantiate an optimizer for both D and G
self.d_optim = optim.Adam(self.D.parameters(), lr=args.lr)
self.g_optim = optim.Adam(self.G.parameters(), lr=args.lr)
def build_mask_discriminator(self):
# Build a discriminator for masks.
D = Discriminator(self.conf.d_mask_params)
D.build()
log.info('Mask Discriminator D_M')
D.model.summary(print_fn=log.info)
self.D_Mask = D.model
real_M = Input(self.conf.d_mask_params.input_shape)
fake_M = Input(self.conf.d_mask_params.input_shape)
real = self.D_Mask(real_M)
fake = self.D_Mask(fake_M)
self.D_Mask_trainer = Model([real_M, fake_M], [real, fake],
name='D_Mask_trainer')
self.D_Mask_trainer.compile(Adam(lr=self.conf.d_mask_params.lr),
loss='mse')
self.D_Mask_trainer.summary(print_fn=log.info)
def build_model(args):
args.to_train = 'CDG'
networks = {}
opts = {}
if 'C' in args.to_train:
networks['C'] = GuidingNet(args.img_size, {'cont': args.sty_dim, 'disc': args.output_k})
networks['C_EMA'] = GuidingNet(args.img_size, {'cont': args.sty_dim, 'disc': args.output_k})
if 'D' in args.to_train:
networks['D'] = Discriminator(args.img_size, num_domains=args.output_k)
if 'G' in args.to_train:
networks['G'] = Generator(args.img_size, args.sty_dim, use_sn=False)
networks['G_EMA'] = Generator(args.img_size, args.sty_dim, use_sn=False)
if args.distributed:
if args.gpu is not None:
print('Distributed to', args.gpu)
torch.cuda.set_device(args.gpu)
args.batch_size = int(args.batch_size / args.ngpus_per_node)
args.workers = int(args.workers / args.ngpus_per_node)
for name, net in networks.items():
if name in ['inceptionNet']:
continue
net_tmp = net.cuda(args.gpu)
networks[name] = torch.nn.parallel.DistributedDataParallel(net_tmp, device_ids=[args.gpu], output_device=args.gpu)
else:
for name, net in networks.items():
net_tmp = net.cuda()
networks[name] = torch.nn.parallel.DistributedDataParallel(net_tmp)
elif args.gpu is not None:
torch.cuda.set_device(args.gpu)
for name, net in networks.items():
networks[name] = net.cuda(args.gpu)
else:
for name, net in networks.items():
networks[name] = torch.nn.DataParallel(net).cuda()
if 'C' in args.to_train:
opts['C'] = torch.optim.Adam(
networks['C'].module.parameters() if args.distributed else networks['C'].parameters(),
1e-4, weight_decay=0.001)
if args.distributed:
networks['C_EMA'].module.load_state_dict(networks['C'].module.state_dict())
else:
networks['C_EMA'].load_state_dict(networks['C'].state_dict())
if 'D' in args.to_train:
opts['D'] = torch.optim.RMSprop(
networks['D'].module.parameters() if args.distributed else networks['D'].parameters(),
1e-4, weight_decay=0.0001)
if 'G' in args.to_train:
opts['G'] = torch.optim.RMSprop(
networks['G'].module.parameters() if args.distributed else networks['G'].parameters(),
1e-4, weight_decay=0.0001)
return networks, opts
def test_model(image):
dvo_feature_extractor = DVOFeature()
dvo_regressor = DVORegression()
dvo_discriminator = Discriminator(500,500,2)
feature = dvo_feature_extractor(image)
print(feature.shape)
motion = dvo_regressor(feature)
print(motion.shape)
dis = dvo_discriminator(feature)
print(dis)
def __init__(self, opt):
super(Model, self).__init__()
# Generator
self.gen = Generator(opt).cuda(opt.gpu_id)
self.gen_params = self.gen.parameters()
num_params = 0
for p in self.gen.parameters():
num_params += p.numel()
print(self.gen)
print(num_params)
# Discriminator
self.dis = Discriminator(opt).cuda(opt.gpu_id)
self.dis_params = self.dis.parameters()
num_params = 0
for p in self.dis.parameters():
num_params += p.numel()
print(self.dis)
print(num_params)
# Regressor
if opt.mse_weight:
self.reg = torch.load('data/utils/classifier.pth').cuda(
opt.gpu_id).eval()
else:
self.reg = None
# Losses
self.criterion_gan = GANLoss(opt, self.dis)
self.criterion_mse = lambda x, y: l1_loss(x, y) * opt.mse_weight
self.loss_mse = Variable(torch.zeros(1).cuda())
self.loss_adv = Variable(torch.zeros(1).cuda())
self.loss = Variable(torch.zeros(1).cuda())
self.path = opt.experiments_dir + opt.experiment_name + '/checkpoints/'
self.gpu_id = opt.gpu_id
self.noise_channels = opt.in_channels - len(opt.input_idx.split(','))
def __init__(self,
n_items,
user_layers,
lambda_mi=0.1,
drop_ratio=0.4,
aggregator_type='attention'):
super(GroupIM, self).__init__()
self.n_items = n_items
self.lambda_mi = lambda_mi
self.drop = nn.Dropout(drop_ratio)
self.embedding_dim = user_layers[-1]
self.aggregator_type = aggregator_type
self.user_preference_encoder = Encoder(self.n_items, user_layers,
self.embedding_dim, drop_ratio)
if self.aggregator_type == 'maxpool':
self.preference_aggregator = MaxPoolAggregator(
self.embedding_dim, self.embedding_dim)
elif self.aggregator_type == 'meanpool':
self.preference_aggregator = MeanPoolAggregator(
self.embedding_dim, self.embedding_dim)
elif self.aggregator_type == 'attention':
self.preference_aggregator = AttentionAggregator(
self.embedding_dim, self.embedding_dim)
else:
raise NotImplementedError(
"Aggregator type {} not implemented ".format(
self.aggregator_type))
self.group_predictor = nn.Linear(self.embedding_dim,
self.n_items,
bias=False)
nn.init.xavier_uniform_(self.group_predictor.weight)
self.discriminator = Discriminator(embedding_dim=self.embedding_dim)
for m in self.modules():
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight)
if isinstance(m, nn.Embedding):
nn.init.xavier_uniform_(m.weight)
def __init__(self, args, training=False):
# Define parameters for training
self.learning_rate = args["learning_rate"] # Learning rate
self.loss_function = args["loss_function"] # loss_function
self.metrics = args["metrics"] # metrics
self.channels = args['channels']
self.batch_size = args['batch_size']
self.num_steps = args['num_steps']
self.seq_shape = (self.num_steps, self.channels)
self.noise_dim = args["noise_dim"] # metrics
self.use_random_noise = args["use_random_noise"]
self.training_mode = training
# Following parameter and optimizer set as recommended in paper
self.optimizer = Adam(lr=self.learning_rate)
# Build and compile the critic
self.critic = Discriminator(args, training)
self.critic.model.compile(loss=self.loss_function,
optimizer=self.optimizer,
metrics=self.metrics)
# Build the generator
self.generator = Generator(args, training)
# Build de combined model (generator => critic)
# The generator takes latent space as input and generates fake signal as output
z = Input(shape=(self.noise_dim, ))
signal = self.generator.model(z)
# For the combined model we will only train the generator
self.critic.model.trainable = False
# The critic takes generated signal as input and determines validity
valid = self.critic.model(signal)
# The combined model (stacked generator and critic) - try to minimize the loss
self.combined = Model(z, valid)
self.combined.compile(loss='binary_crossentropy',
optimizer=self.optimizer)
def __init__(self, hparams):
super().__init__()
self.hparams = hparams
# Encoder
self.encoder = Encoder(ngf=self.hparams.ngf, z_dim=self.hparams.z_dim)
self.encoder.apply(weights_init)
device = "cuda" if isinstance(self.hparams.gpus, int) else "cpu"
# Decoder
self.decoder = Decoder(ngf=self.hparams.ngf, z_dim=self.hparams.z_dim)
self.decoder.apply(weights_init)
# Discriminator
self.discriminator = Discriminator()
self.discriminator.apply(weights_init)
# Losses
self.criterionFeat = torch.nn.L1Loss()
self.criterionGAN = GANLoss(gan_mode="lsgan")
if self.hparams.use_vgg:
self.criterion_perceptual_style = [Perceptual_Loss(device)]
def test_test_real():
args = parse()
print(args)
dataset = SepeDataset(args.poses_train,args.images_train,coor_layer_flag =False)
dataloader = DataLoader(dataset, batch_size=10,shuffle=False ,num_workers=1,drop_last=True)
dvo_feature_extractor = DVOFeature()
dvo_regressor = DVORegression()
dvo_discriminator = Discriminator(500,500,2)
dvo_feature_extractor.load_state_dict(torch.load(('feature_ntsd_2_10.pt')))
dvo_regressor.load_state_dict(torch.load('regressor_seed_ntsd_2_10.pt'))
test(dvo_feature_extractor,dvo_regressor,dataloader,args)
def test_train_real():
args = parse()
print(args)
dataset = SepeDataset(args.poses_train,args.images_train,coor_layer_flag =False)
dataloader = DataLoader(dataset, batch_size=3,shuffle=True ,num_workers=1,drop_last=True,worker_init_fn=lambda wid:np.random.seed(np.uint32(torch.initial_seed() + wid)))
dvo_feature_extractor = DVOFeature()
dvo_regressor = DVORegression()
dvo_discriminator = Discriminator(500,500,2)
trained_feature,trained_regressor = train(dvo_feature_extractor,dvo_regressor,dataloader,args)
torch.save(trained_feature.state_dict(),'feature_'+args.tag+str(args.epoch)+'.pt')
torch.save(trained_regressor.state_dict(),'regressor_'+args.tag+str(args.epoch)+'.pt')
def test_train():
args = parse()
print(args)
motion_ax_i = [int(i) for i in args.motion_ax.split(' ')]
dataset = RandomDataset(20000,motion_ax = motion_ax_i)
dataloader = DataLoader(dataset, batch_size=1000,shuffle=False ,num_workers=1,drop_last=True,worker_init_fn=lambda wid:np.random.seed(np.uint32(torch.initial_seed() + wid)))
dvo_feature_extractor = DVOFeature()
dvo_regressor = DVORegression()
dvo_discriminator = Discriminator(500,500,2)
trained_feature,trained_regressor = train(dvo_feature_extractor,dvo_regressor,dataloader,args)
torch.save(trained_feature.state_dict(),'feature_seed'+args.motion_ax.replace(' ','')+str(args.epoch)+'.pt')
torch.save(trained_regressor.state_dict(),'regressor_seed'+args.motion_ax.replace(' ','')+str(args.epoch)+'.pt')
def build_image_discriminator2(self):
"""
Build a discriminator for images
"""
params2 = self.conf.d_image_params
params2['name'] = 'D_Image2'
D = Discriminator(params2)
D.build()
log.info('Image Discriminator D_I2')
D.model.summary(print_fn=log.info)
self.D_Image2 = D.model
real_x = Input(self.conf.d_image_params.input_shape)
fake_x = Input(self.conf.d_image_params.input_shape)
real = self.D_Image2(real_x)
fake = self.D_Image2(fake_x)
self.D_Image2_trainer = Model([real_x, fake_x], [real, fake],
name='D_Image2_trainer')
self.D_Image2_trainer.compile(Adam(lr=self.conf.d_image_params.lr),
loss='mse')
self.D_Image2_trainer.summary(print_fn=log.info)
def test_test():
args = parse()
print(args)
motion_ax_i = [int(i) for i in args.motion_ax.split(' ')]
test_motion_ax_i = [int(i) for i in args.test_motion_ax.split(' ')]
dataset = RandomDataset(2,motion_ax = test_motion_ax_i)
dataloader = DataLoader(dataset, batch_size=1,shuffle=False ,num_workers=1,drop_last=True)
dvo_feature_extractor = DVOFeature()
dvo_regressor = DVORegression()
dvo_discriminator = Discriminator(500,500,2)
dvo_feature_extractor.load_state_dict(torch.load('feature'+args.motion_ax.replace(' ','')+'.pt'))
dvo_regressor.load_state_dict(torch.load('regressor'+args.motion_ax.replace(' ','')+'.pt'))
test(dvo_feature_extractor,dvo_regressor,dataloader,args)
def build_models(conf, n_voc, seq_len):
model = dict()
netG = Generator(n_voc, conf.emb_dim, conf.hid_dim, seq_len, conf.gpu_num)
netD = Discriminator(n_voc, conf.emb_dim, conf.hid_dim, seq_len,
conf.gpu_num)
if conf.gpu_num == 0:
model["netG_0"] = netG
model["netD_0"] = netD
else:
for i in range(conf.gpu_num):
copy_netG = netG.copy()
copy_netD = netD.copy()
copy_netG.gpu_id = i
copy_netD.gpu_id = i
model[f"netG_{i}"] = copy_netG.to_gpu(i)
model[f"netD_{i}"] = copy_netD.to_gpu(i)
return model
def main():
data_loader = DataLoader(data_path, batch_size)
generator = Generator(vocab_size, g_emb_dim, g_hidden_dim)
discriminator = Discriminator(vocab_size, d_hidden_dim)
gen_optimizer = optim.Adam(generator.parameters(), lr=0.0001)
disc_optimizer = optim.Adam(discriminator.parameters(), lr=0.0001)
bce_criterion = nn.BCELoss()
gen_criterion = nn.NLLLoss(size_average=False)
if (opt.cuda):
generator.cuda()
pretrain_lstm(generator, data_loader, gen_optimizer, gen_criterion, 10)
all_G_losses = []
all_D_losses = []
for i in tqdm(range(total_epochs)):
g_losses, d_losses = train_gan_epoch(discriminator, generator,
data_loader, gen_optimizer,
disc_optimizer, bce_criterion)
all_G_losses += g_losses
all_D_losses += d_losses
sample = generator.sample(batch_size, g_seq_length)
print(generator)
with open('./data/gumbel_softmax_gan_gen.txt', 'w') as f:
for each_str in data_loader.convert_to_char(sample):
f.write(each_str + '\n')
gen_file_name = 'gen_gumbel_softmax_' + str(total_epochs) + '.pth'
disc_file_name = 'disc_gumbel_softmax_' + str(total_epochs) + '.pth'
Utils.save_checkpoints(checkpoint_dir, gen_file_name, generator)
Utils.save_checkpoints(checkpoint_dir, disc_file_name, discriminator)
plt.plot(list(range(len(all_G_losses))),
all_G_losses,
'g-',
label='gen loss')
plt.plot(list(range(len(all_D_losses))),
all_D_losses,
'b-',
label='disc loss')
plt.legend()
plt.show()
def __init__(self, config):
"""
Construct a new GAN trainer
:param Config config: The parsed network configuration.
"""
self.config = config
LOG.info("CUDA version: {0}".format(version.cuda))
LOG.info("Creating data loader from path {0}".format(config.FILENAME))
self.data_loader = Data(
config.FILENAME,
config.BATCH_SIZE,
polarisations=config.POLARISATIONS, # Polarisations to use
frequencies=config.FREQUENCIES, # Frequencies to use
max_inputs=config.
MAX_SAMPLES, # Max inputs per polarisation and frequency
normalise=config.NORMALISE) # Normalise inputs
shape = self.data_loader.get_input_shape()
width = shape[1]
LOG.info("Creating models with input shape {0}".format(shape))
self._autoencoder = Autoencoder(width)
self._discriminator = Discriminator(width)
# TODO: Get correct input and output widths for generator
self._generator = Generator(width, width)
if config.USE_CUDA:
LOG.info("Using CUDA")
self.autoencoder = self._autoencoder.cuda()
self.discriminator = self._discriminator.cuda()
self.generator = self._generator.cuda()
else:
LOG.info("Using CPU")
self.autoencoder = self._autoencoder
self.discriminator = self._discriminator
self.generator = self._generator
def init_swa_models(self):
"""
Initialise objects for Stochastic Weight Averaging
"""
self.swa_D_Mask = SWA(40,
Discriminator(self.conf.d_mask_params).build,
None)
self.swa_D_Image1 = SWA(40,
Discriminator(self.conf.d_image_params).build,
None)
self.swa_D_Image2 = SWA(40,
Discriminator(self.conf.d_image_params).build,
None)
self.swa_Enc_Anatomy1 = SWA(40, anatomy_encoder.build,
self.conf.anatomy_encoder)
self.swa_Enc_Anatomy2 = SWA(40, anatomy_encoder.build,
self.conf.anatomy_encoder)
self.swa_Enc_Modality = SWA(40, modality_encoder.build, self.conf)
self.swa_Anatomy_Fuser = SWA(40, anatomy_fuser.build, self.conf)
self.swa_Segmentor = SWA(40, segmentor.build, self.conf)
self.swa_Decoder = SWA(40, decoder.build, self.conf)
self.swa_Balancer = SWA(40, balancer.build, self.conf)
self.set_swa_model_weights()
def __init__(self, device, model, image_nc, box_min, box_max):
output_nc = image_nc
self.device = device
self.model = model
self.input_nc = image_nc
self.output_nc = output_nc
self.box_min = box_min
self.box_max = box_max
self.gen_input_nc = image_nc
self.netG = Generator(self.gen_input_nc, image_nc).to(device)
self.netDisc = Discriminator(image_nc).to(device)
# initialize all weights
self.netG.apply(weights_init)
self.netDisc.apply(weights_init)
# initialize optimizers
self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=0.001)
self.optimizer_D = torch.optim.Adam(self.netDisc.parameters(), lr=0.001)
if not os.path.exists(models_path):
os.makedirs(models_path)
def build_model(args):
generator = Generator(args)
discriminator = Discriminator(args)
mapping_network = MappingNetwork(args)
generator_ema = copy.deepcopy(generator)
mapping_network_ema = copy.deepcopy(mapping_network)
nets = Munch(generator=generator,
discriminator=discriminator,
mapping_network=mapping_network)
nets_ema = Munch(generator=generator_ema,
mapping_network=mapping_network_ema)
return nets, nets_ema
