def model_init(self):
self.D = Discriminator(self.model_type, self.image_size,
self.hidden_dim, self.n_filter, self.n_repeat)
self.G = Generator(self.model_type, self.image_size, self.hidden_dim,
self.n_filter, self.n_repeat)
self.D = cuda(self.D, self.cuda)
self.G = cuda(self.G, self.cuda)
self.D.weight_init(mean=0.0, std=0.02)
self.G.weight_init(mean=0.0, std=0.02)
self.D_optim = optim.Adam(self.D.parameters(),
lr=self.D_lr,
betas=(0.5, 0.999))
self.G_optim = optim.Adam(self.G.parameters(),
lr=self.G_lr,
betas=(0.5, 0.999))
#self.D_optim_scheduler = lr_scheduler.ExponentialLR(self.D_optim, gamma=0.97)
#self.G_optim_scheduler = lr_scheduler.ExponentialLR(self.G_optim, gamma=0.97)
self.D_optim_scheduler = lr_scheduler.StepLR(self.D_optim,
step_size=1,
gamma=0.5)
self.G_optim_scheduler = lr_scheduler.StepLR(self.G_optim,
step_size=1,
gamma=0.5)
if not self.ckpt_dir.exists():
self.ckpt_dir.mkdir(parents=True, exist_ok=True)
if self.load_ckpt:
self.load_checkpoint()
def load_model(model_file, hidden_size, upsampling, cuda=False):
if cuda:
from_before = torch.load(model_file)
else:
from_before = torch.load(model_file,
map_location=lambda storage, loc: storage)
total_examples = from_before['total_examples']
gen_losses = from_before['gen_losses']
disc_losses = from_before['disc_losses']
gen_loss_per_epoch = from_before['gen_loss_per_epoch']
disc_loss_per_epoch = from_before['disc_loss_per_epoch']
gen_state_dict = from_before['gen_state_dict']
disc_state_dict = from_before['disc_state_dict']
fixed_noise = from_before['fixed_noise']
epoch = from_before['epoch']
# load generator and discriminator
if upsampling == 'transpose':
from models.model import Generator, Discriminator
elif upsampling == 'nn':
from models.model_nn import Generator, Discriminator
elif upsampling == 'bilinear':
from models.model_bilinear import Generator, Discriminator
gen = Generator(hidden_dim=hidden_size,
dropout=0.4) # TODO: save dropout in checkpoint
disc = Discriminator(leaky=0.2, dropout=0.4) # TODO: same here
disc.load_state_dict(disc_state_dict)
gen.load_state_dict(gen_state_dict)
return total_examples, fixed_noise, gen_losses, disc_losses, \
gen_loss_per_epoch, disc_loss_per_epoch, epoch, gen, disc
def main():
if not osp.isdir(args.checkpoint):
mkdir_p(args.checkpoint)
if torch.cuda.is_available():
device = torch.device('cuda:0')
else:
device = torch.device('cpu')
L = Lifter().to(device)
D = Discriminator().to(device)
T = Discriminator().to(device)
optim_L = optim.Adam(L.parameters(), lr=args.lift_lr)
optim_D = optim.Adam(D.parameters(), lr=args.disc_lr)
optim_T = optim.Adam(T.parameters(), lr=args.disc_lr)
# use 2D results from Stack Hourglass Net
train_loader = data.DataLoader(
H36M(length=args.length, action='all', is_train=True, use_sh_detection=True),
batch_size=1024,
shuffle=True,
pin_memory=True,
)
test_loader = data.DataLoader(
H36M(length=1, action='all', is_train=False, use_sh_detection=True),
batch_size=512,
shuffle=False,
)
# Logger
logger = Logger(osp.join(args.checkpoint, 'log.txt'), title='Human3.6M')
logger_err = Logger(osp.join(args.checkpoint, 'log_err.txt'), title='Human3.6M MPJPE err')
logger.set_names(['2d_loss ', '3d_loss ', 'adv_loss ', 'temporal_loss '])
logger_err.set_names(['err'])
for epoch in range(args.epoches):
print('\nEpoch: [%d / %d]' % (epoch+1, args.epoches))
loss_2d, loss_3d, loss_adv, loss_t = train(train_loader, L, D, T, optim_L, optim_D, optim_T, epoch+1, device, args)
logger.append([loss_2d, loss_3d, loss_adv, loss_t])
if (epoch + 1) % args.checkpoint_save_interval == 0:
save_checkpoint({
'epoch': epoch + 1,
'state_dict_L': L.state_dict(),
'state_dict_D': D.state_dict(),
'state_dict_T': T.state_dict(),
}, checkpoint=args.checkpoint)
if (epoch + 1) % args.eval_interval == 0:
ttl_err = test(test_loader, L, epoch, device, args)
logger_err.append([ttl_err])
logger.close()
logger_err.close()
def train():
args = parse_args()
cfg = Config.from_file(args.config)
# Dimensionality of the latent vector.
latent_size = cfg.models.generator.z_dim
# Use sigmoid activation for the last layer?
cfg.models.discriminator.sigmoid_at_end = cfg.train.loss_type in [
'ls', 'gan'
]
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
G = Generator(model_cfg=cfg.models.generator,
target_size=cfg.train.target_size)
D = Discriminator(model_cfg=cfg.models.discriminator,
target_size=cfg.train.target_size)
#print(G)
#print(D)
dataset = FaceDataset(cfg.train.dataset)
assert len(dataset) > 0
print(f'train dataset contains {len(dataset)} images.')
clip = cfg.models.generator.z_clipping if hasattr(cfg.models.generator,
'z_zlipping') else None
z_generator = RandomNoiseGenerator(cfg.models.generator.z_dim,
'gaussian',
clip=clip)
pggan = PGGAN(G, D, dataset, z_generator, args.gpu, cfg, args.resume)
pggan.train()
def load_model(opt):
print('initializing model ...')
print("LOADING GENERATOR MODEL")
model_g = Generator(100, gpu=opt.SYSTEM.USE_GPU)
print("LOADING DISCRIMINATOR MODEL")
model_d = Discriminator(output_dim=opt.TRAIN.TOTAL_FEATURES,
gpu=opt.SYSTEM.USE_GPU)
return model_g, model_d
def test_discriminator_shape(self):
disc = Discriminator(img_size=self.img_size)
img = torch.randn([1, 3, 1024, 1024])
logits_shape = (1, 1, 5, 5)
decode_img_part = (1, 3, 128, 128)
decoded_img = (1, 3, 128, 128)
real_fake_logits_out, decoded_img_128_part, decoded_img_128 = disc(img)
self.assertEqual(real_fake_logits_out.shape, logits_shape)
self.assertEqual(decoded_img_128_part.shape, decoded_img_128.shape)
self.assertEqual(decoded_img_128_part.shape, decode_img_part)
self.assertEqual(decoded_img_128.shape, decoded_img)
help='rampup_kimg.')
parser.add_argument('--rampdown_kimg',
default=10000,
type=float,
help='rampdown_kimg.')
# TODO: support conditional inputs
args = parser.parse_args()
opts = {k: v for k, v in args._get_kwargs()}
latent_size = 512
sigmoid_at_end = args.gan in ['lsgan', 'gan']
G = Generator(num_channels=3,
latent_size=latent_size,
resolution=args.target_resol,
fmap_max=latent_size,
fmap_base=8192,
tanh_at_end=False)
D = Discriminator(num_channels=3,
resolution=args.target_resol,
fmap_max=latent_size,
fmap_base=8192,
sigmoid_at_end=sigmoid_at_end)
print(G)
print(D)
data = CelebA()
noise = RandomNoiseGenerator(latent_size, 'gaussian')
pggan = PGGAN(G, D, data, noise, opts)
pggan.train()
train_loader, valid_loader = make_dataloader(opt)
# Decide which device we want to run on
device = torch.device("cuda:0" if (
torch.cuda.is_available() and opt.ngpu > 0) else "cpu")
# Create the generator
netG = Generator().to(device)
if (device.type == 'cuda') and (opt.ngpu > 1):
netG = nn.DataParallel(netG, list(range(opt.ngpu)))
netG.apply(weights_init)
print(netG)
# Create the Discriminator
netD = Discriminator().to(device)
if (device.type == 'cuda') and (opt.ngpu > 1):
netD = nn.DataParallel(netD, list(range(opt.ngpu)))
netD.apply(weights_init)
print(netD)
# Initialize BCELoss function
criterion = nn.BCELoss()
l1_loss = nn.L1Loss(reduction='sum')
# Establish convention for real and fake labels during training
real_label = 1
fake_label = 0
# Setup Adam optimizers for both G and D
optimizerD = optim.Adam(netD.parameters(), lr=0.0002, betas=(0.5, 0.9))
def main(train_set, learning_rate, n_epochs, beta_0, beta_1, batch_size,
num_workers, hidden_size, model_file, cuda, display_result_every,
checkpoint_interval, seed, label_smoothing, grad_clip, dropout,
upsampling):
# make data between -1 and 1
data_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
train_dataset = datasets.ImageFolder(root=os.path.join(
os.getcwd(), train_set),
transform=data_transform)
train_dataloader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
drop_last=True)
# initialize model
if model_file:
try:
total_examples, fixed_noise, gen_losses, disc_losses, gen_loss_per_epoch, \
disc_loss_per_epoch, prev_epoch, gen, disc = load_model(model_file, hidden_size, upsampling, cuda)
print('model loaded successfully!')
except:
print('could not load model! creating new model...')
model_file = None
if not model_file:
print('creating new model...')
if upsampling == 'transpose':
from models.model import Generator, Discriminator
elif upsampling == 'nn':
from models.model_nn import Generator, Discriminator
elif upsampling == 'bilinear':
from models.model_bilinear import Generator, Discriminator
gen = Generator(hidden_dim=hidden_size, leaky=0.2, dropout=dropout)
disc = Discriminator(leaky=0.2, dropout=dropout)
gen.weight_init(mean=0, std=0.02)
disc.weight_init(mean=0, std=0.02)
total_examples = 0
disc_losses = []
gen_losses = []
disc_loss_per_epoch = []
gen_loss_per_epoch = []
prev_epoch = 0
# Sample minibatch of m noise samples from noise prior p_g(z) and transform
if cuda:
fixed_noise = Variable(torch.randn(9, hidden_size).cuda())
else:
fixed_noise = Variable(torch.rand(9, hidden_size))
if cuda:
gen.cuda()
disc.cuda()
# Binary Cross Entropy loss
BCE_loss = nn.BCELoss()
# Adam optimizer
gen_optimizer = optim.Adam(gen.parameters(),
lr=learning_rate,
betas=(beta_0, beta_1),
eps=1e-8)
disc_optimizer = optim.Adam(disc.parameters(),
lr=learning_rate,
betas=(beta_0, beta_1),
eps=1e-8)
# results save folder
gen_images_dir = 'results/generated_images'
train_summaries_dir = 'results/training_summaries'
checkpoint_dir = 'results/checkpoints'
if not os.path.isdir('results'):
os.mkdir('results')
if not os.path.isdir(gen_images_dir):
os.mkdir(gen_images_dir)
if not os.path.isdir(train_summaries_dir):
os.mkdir(train_summaries_dir)
if not os.path.isdir(checkpoint_dir):
os.mkdir(checkpoint_dir)
np.random.seed(
seed
) # reset training seed to ensure that batches remain the same between runs!
try:
for epoch in range(prev_epoch, n_epochs):
disc_losses_epoch = []
gen_losses_epoch = []
for idx, (true_batch, _) in enumerate(train_dataloader):
disc.zero_grad()
# hack 6 of https://github.com/soumith/ganhacks
if label_smoothing:
true_target = torch.FloatTensor(batch_size).uniform_(
0.7, 1.2)
else:
true_target = torch.ones(batch_size)
# Sample minibatch of examples from data generating distribution
if cuda:
true_batch = Variable(true_batch.cuda())
true_target = Variable(true_target.cuda())
else:
true_batch = Variable(true_batch)
true_target = Variable(true_target)
# train discriminator on true data
true_disc_result = disc.forward(true_batch)
disc_train_loss_true = BCE_loss(true_disc_result.squeeze(),
true_target)
disc_train_loss_true.backward()
torch.nn.utils.clip_grad_norm(disc.parameters(), grad_clip)
# Sample minibatch of m noise samples from noise prior p_g(z) and transform
if label_smoothing:
fake_target = torch.FloatTensor(batch_size).uniform_(
0, 0.3)
else:
fake_target = torch.zeros(batch_size)
if cuda:
z = Variable(torch.randn(batch_size, hidden_size).cuda())
fake_target = Variable(fake_target.cuda())
else:
z = Variable(torch.randn(batch_size, hidden_size))
fake_target = Variable(fake_target)
# train discriminator on fake data
fake_batch = gen.forward(z.view(-1, hidden_size, 1, 1))
fake_disc_result = disc.forward(fake_batch.detach(
)) # detach so gradients not computed for generator
disc_train_loss_false = BCE_loss(fake_disc_result.squeeze(),
fake_target)
disc_train_loss_false.backward()
torch.nn.utils.clip_grad_norm(disc.parameters(), grad_clip)
disc_optimizer.step()
# compute performance statistics
disc_train_loss = disc_train_loss_true + disc_train_loss_false
disc_losses_epoch.append(disc_train_loss.data[0])
disc_fake_accuracy = 1 - fake_disc_result.mean().data[0]
disc_true_accuracy = true_disc_result.mean().data[0]
# Sample minibatch of m noise samples from noise prior p_g(z) and transform
if label_smoothing:
true_target = torch.FloatTensor(batch_size).uniform_(
0.7, 1.2)
else:
true_target = torch.ones(batch_size)
if cuda:
z = Variable(torch.randn(batch_size, hidden_size).cuda())
true_target = Variable(true_target.cuda())
else:
z = Variable(torch.rand(batch_size, hidden_size))
true_target = Variable(true_target)
# train generator
gen.zero_grad()
fake_batch = gen.forward(z.view(-1, hidden_size, 1, 1))
disc_result = disc.forward(fake_batch)
gen_train_loss = BCE_loss(disc_result.squeeze(), true_target)
gen_train_loss.backward()
torch.nn.utils.clip_grad_norm(gen.parameters(), grad_clip)
gen_optimizer.step()
gen_losses_epoch.append(gen_train_loss.data[0])
if (total_examples != 0) and (total_examples %
display_result_every == 0):
print(
'epoch {}: step {}/{} disc true acc: {:.4f} disc fake acc: {:.4f} '
'disc loss: {:.4f}, gen loss: {:.4f}'.format(
epoch + 1, idx + 1, len(train_dataloader),
disc_true_accuracy, disc_fake_accuracy,
disc_train_loss.data[0], gen_train_loss.data[0]))
# Checkpoint model
total_examples += batch_size
if (total_examples != 0) and (total_examples %
checkpoint_interval == 0):
disc_losses.extend(disc_losses_epoch)
gen_losses.extend(gen_losses_epoch)
save_checkpoint(total_examples=total_examples,
fixed_noise=fixed_noise,
disc=disc,
gen=gen,
gen_losses=gen_losses,
disc_losses=disc_losses,
disc_loss_per_epoch=disc_loss_per_epoch,
gen_loss_per_epoch=gen_loss_per_epoch,
epoch=epoch,
directory=checkpoint_dir)
print("Checkpoint saved!")
# sample images for inspection
save_image_sample(batch=gen.forward(
fixed_noise.view(-1, hidden_size, 1, 1)),
cuda=cuda,
total_examples=total_examples,
directory=gen_images_dir)
print("Saved images!")
# save learning curves for inspection
save_learning_curve(gen_losses=gen_losses,
disc_losses=disc_losses,
total_examples=total_examples,
directory=train_summaries_dir)
print("Saved learning curves!")
disc_loss_per_epoch.append(np.average(disc_losses_epoch))
gen_loss_per_epoch.append(np.average(gen_losses_epoch))
# Save epoch learning curve
save_learning_curve_epoch(gen_losses=gen_loss_per_epoch,
disc_losses=disc_loss_per_epoch,
total_epochs=epoch + 1,
directory=train_summaries_dir)
print("Saved learning curves!")
print('epoch {}/{} disc loss: {:.4f}, gen loss: {:.4f}'.format(
epoch + 1, n_epochs,
np.array(disc_losses_epoch).mean(),
np.array(gen_losses_epoch).mean()))
disc_losses.extend(disc_losses_epoch)
gen_losses.extend(gen_losses_epoch)
except KeyboardInterrupt:
print("Saving before quit...")
save_checkpoint(total_examples=total_examples,
fixed_noise=fixed_noise,
disc=disc,
gen=gen,
disc_loss_per_epoch=disc_loss_per_epoch,
gen_loss_per_epoch=gen_loss_per_epoch,
gen_losses=gen_losses,
disc_losses=disc_losses,
epoch=epoch,
directory=checkpoint_dir)
print("Checkpoint saved!")
# sample images for inspection
save_image_sample(batch=gen.forward(
fixed_noise.view(-1, hidden_size, 1, 1)),
cuda=cuda,
total_examples=total_examples,
directory=gen_images_dir)
print("Saved images!")
# save learning curves for inspection
save_learning_curve(gen_losses=gen_losses,
disc_losses=disc_losses,
total_examples=total_examples,
directory=train_summaries_dir)
print("Saved learning curves!")
batch_size = args.batch_size
num_epochs = args.max_epoch
input_dim = 1
hidden_dim = 32
output_dim = 1
num_epochs = 100000
num_epochs_pre = 500
learning_rate = 0.03
# Samples
data = DataDistribution(mu, sigma)
gen = NoiseDistribution(data_range)
# Models
G = Generator(input_dim, hidden_dim, output_dim)
D = Discriminator(input_dim, hidden_dim, output_dim)
# Loss function
criterion = torch.nn.BCELoss()
# optimizer
optimizer = torch.optim.SGD(D.parameters(), lr=learning_rate)
D_pre_losses = []
num_samples_pre = 5000
num_bins_pre = 100
for epoch in range(num_epochs_pre):
# Generate samples
d = data.sample(num_samples_pre)
histc, edges = np.histogram(d, num_bins_pre, density=True)
print(f"=> Called with args {args.__dict__}")
print(f"=> Config params {cfg.__dict__}")
print(f"=> Run on device {device}")
# define dataset and dataloader
dataset = AnimeFacesDataset(args.data_path)
cfg.DATASET_SIZE = len(dataset)
dataloader = DataLoader(dataset,
batch_size=cfg.BATCH_SIZE,
shuffle=True,
num_workers=2,
drop_last=True,
pin_memory=True)
# define models
gen = Generator(cfg.Z_DIMENSION, cfg.CHANNELS_IMG,
cfg.FEATURES_GEN).to(device)
disc = Discriminator(cfg.CHANNELS_IMG, cfg.FEATURES_DISC).to(device)
if args.checkpoint_path:
opt_gen = optim.Adam(gen.parameters(),
lr=cfg.LEARNING_RATE,
betas=(0.5, 0.999))
opt_disc = optim.Adam(disc.parameters(),
lr=cfg.LEARNING_RATE,
betas=(0.5, 0.999))
cp = torch.load(args.checkpoint_path)
start_epoch, end_epoch, fixed_noise = load_checkpoint(
cp, gen, disc, opt_gen, opt_disc)
cfg.NUM_EPOCHS = end_epoch
else:
print("=> Init default weights of models and fixed noise")
# FIXME sometime (usually) when the weights is initialized from normal distribution can cause mode collapse
test_set = MVTec_dataset(test_img_paths,
test_labels,
transform=transform_mvtec_test)
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=0)
test_loader = torch.utils.data.DataLoader(test_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=0)
#-------------------------------------
D = Discriminator()
model = VAE(args)
resN = ResNet_(args)
model = model.to(device)
D = D.to(device)
resN = resN.to(device)
metric_fc = AdaCos(num_features=args.num_features,
num_classes=args.num_classes).to(device)
criterion = nn.CrossEntropyLoss().to(device)
cudnn.benchmark = True
optimizer = optim.SGD(filter(lambda p: p.requires_grad, resN.parameters()),
lr=args.lr,
momentum=args.momentum,
import torch
from models.model import Encoder
from models.model import Decoder
from models.model import Discriminator
if __name__ == "__main__":
encoder = Encoder()
img = torch.FloatTensor(1, 1, 32, 32)
output = encoder(img)
print(output.size())
decoder = Decoder()
img = decoder(output)
print(img.size())
discriminator = Discriminator()
output = discriminator(img)
print(output.size())
class BEGAN(object):
def __init__(self, args):
# Misc
self.args = args
self.cuda = args.cuda and torch.cuda.is_available()
self.sample_num = 100
# Optimization
self.epoch = args.epoch
self.batch_size = args.batch_size
self.D_lr = args.D_lr
self.G_lr = args.G_lr
self.gamma = args.gamma
self.lambda_k = args.lambda_k
self.Kt = 0.0
self.global_epoch = 0
self.global_iter = 0
# Visualization
self.env_name = args.env_name
self.visdom = args.visdom
self.port = args.port
self.timestep = args.timestep
self.output_dir = Path(args.output_dir).joinpath(args.env_name)
self.visualization_init()
# Network
self.model_type = args.model_type
self.n_filter = args.n_filter
self.n_repeat = args.n_repeat
self.image_size = args.image_size
self.hidden_dim = args.hidden_dim
self.fixed_z = Variable(cuda(self.sample_z(self.sample_num),
self.cuda))
self.ckpt_dir = Path(args.ckpt_dir).joinpath(args.env_name)
self.load_ckpt = args.load_ckpt
self.model_init()
# Dataset
self.dataset = args.dataset
self.data_loader = return_data(args)
self.lr_step_size = len(self.data_loader['train'].dataset
) // self.batch_size * self.epoch // 8
def model_init(self):
self.D = Discriminator(self.model_type, self.image_size,
self.hidden_dim, self.n_filter, self.n_repeat)
self.G = Generator(self.model_type, self.image_size, self.hidden_dim,
self.n_filter, self.n_repeat)
self.D = cuda(self.D, self.cuda)
self.G = cuda(self.G, self.cuda)
self.D.weight_init(mean=0.0, std=0.02)
self.G.weight_init(mean=0.0, std=0.02)
self.D_optim = optim.Adam(self.D.parameters(),
lr=self.D_lr,
betas=(0.5, 0.999))
self.G_optim = optim.Adam(self.G.parameters(),
lr=self.G_lr,
betas=(0.5, 0.999))
#self.D_optim_scheduler = lr_scheduler.ExponentialLR(self.D_optim, gamma=0.97)
#self.G_optim_scheduler = lr_scheduler.ExponentialLR(self.G_optim, gamma=0.97)
self.D_optim_scheduler = lr_scheduler.StepLR(self.D_optim,
step_size=1,
gamma=0.5)
self.G_optim_scheduler = lr_scheduler.StepLR(self.G_optim,
step_size=1,
gamma=0.5)
if not self.ckpt_dir.exists():
self.ckpt_dir.mkdir(parents=True, exist_ok=True)
if self.load_ckpt:
self.load_checkpoint()
def visualization_init(self):
if not self.output_dir.exists():
self.output_dir.mkdir(parents=True, exist_ok=True)
if self.visdom:
self.viz_train_curves = visdom.Visdom(env=self.env_name +
'/train_curves',
port=self.port)
self.viz_train_samples = visdom.Visdom(env=self.env_name +
'/train_samples',
port=self.port)
self.viz_test_samples = visdom.Visdom(env=self.env_name +
'/test_samples',
port=self.port)
self.viz_interpolations = visdom.Visdom(env=self.env_name +
'/interpolations',
port=self.port)
self.win_moc = None
def sample_z(self, batch_size=0, dim=0, dist='uniform'):
if batch_size == 0:
batch_size = self.batch_size
if dim == 0:
dim = self.hidden_dim
if dist == 'normal':
return torch.randn(batch_size, dim)
elif dist == 'uniform':
return torch.rand(batch_size, dim).mul(2).add(-1)
else:
return None
def sample_img(self, _type='fixed', nrow=10):
self.set_mode('eval')
if _type == 'fixed':
z = self.fixed_z
elif _type == 'random':
z = self.sample_z(self.sample_num)
z = Variable(cuda(z, self.cuda))
else:
self.set_mode('train')
return
samples = self.unscale(self.G(z))
samples = samples.data.cpu()
filename = self.output_dir.joinpath(_type + ':' +
str(self.global_iter) + '.jpg')
grid = make_grid(samples, nrow=nrow, padding=2, normalize=False)
save_image(grid, filename=filename)
if self.visdom:
self.viz_test_samples.image(grid,
opts=dict(title=str(filename),
nrow=nrow,
factor=2))
self.set_mode('train')
return grid
def set_mode(self, mode='train'):
if mode == 'train':
self.G.train()
self.D.train()
elif mode == 'eval':
self.G.eval()
self.D.eval()
else:
raise ('mode error. It should be either train or eval')
def scheduler_step(self):
self.D_optim_scheduler.step()
self.G_optim_scheduler.step()
def unscale(self, tensor):
return tensor.mul(0.5).add(0.5)
def save_checkpoint(self, filename='ckpt.tar'):
model_states = {'G': self.G.state_dict(), 'D': self.D.state_dict()}
optim_states = {
'G_optim': self.G_optim.state_dict(),
'D_optim': self.D_optim.state_dict()
}
states = {
'iter': self.global_iter,
'epoch': self.global_epoch,
'args': self.args,
'win_moc': self.win_moc,
'fixed_z': self.fixed_z.data.cpu(),
'model_states': model_states,
'optim_states': optim_states
}
file_path = self.ckpt_dir.joinpath(filename)
torch.save(states, file_path.open('wb+'))
print("=> saved checkpoint '{}' (iter {})".format(
file_path, self.global_iter))
def load_checkpoint(self, filename='ckpt.tar'):
file_path = self.ckpt_dir.joinpath(filename)
if file_path.is_file():
checkpoint = torch.load(file_path.open('rb'))
self.global_iter = checkpoint['iter']
self.global_epoch = checkpoint['epoch']
self.win_moc = checkpoint['win_moc']
self.fixed_z = checkpoint['fixed_z']
self.fixed_z = Variable(cuda(self.fixed_z, self.cuda))
self.G.load_state_dict(checkpoint['model_states']['G'])
self.D.load_state_dict(checkpoint['model_states']['D'])
self.G_optim.load_state_dict(checkpoint['optim_states']['G_optim'])
self.D_optim.load_state_dict(checkpoint['optim_states']['D_optim'])
print("=> loaded checkpoint '{} (iter {})'".format(
file_path, self.global_iter))
else:
print("=> no checkpoint found at '{}'".format(file_path))
def train(self):
self.set_mode('train')
for e in range(self.epoch):
self.global_epoch += 1
e_elapsed = time.time()
for idx, (images, labels) in enumerate(self.data_loader['train']):
self.global_iter += 1
# Discriminator Training
x_real = Variable(cuda(images, self.cuda))
D_real = self.D(x_real)
D_loss_real = F.l1_loss(D_real, x_real)
z = self.sample_z()
z = Variable(cuda(z, self.cuda))
x_fake = self.G(z)
D_fake = self.D(x_fake.detach())
D_loss_fake = F.l1_loss(D_fake, x_fake)
D_loss = D_loss_real - self.Kt * D_loss_fake
self.D_optim.zero_grad()
D_loss.backward()
self.D_optim.step()
# Generator Training
z = self.sample_z()
z = Variable(cuda(z, self.cuda))
x_fake = self.G(z)
D_fake = self.D(x_fake)
G_loss = F.l1_loss(x_fake, D_fake)
self.G_optim.zero_grad()
G_loss.backward()
self.G_optim.step()
# Kt update
balance = (self.gamma * D_loss_real - D_loss_fake).data[0]
self.Kt = max(min(self.Kt + self.lambda_k * balance, 1.0), 0.0)
# Visualize process
if self.visdom and self.global_iter % 1000 == 0:
self.viz_train_samples.images(
self.unscale(x_fake).data.cpu(),
opts=dict(
title='x_fake:{:d}'.format(self.global_iter)))
self.viz_train_samples.images(
self.unscale(D_fake).data.cpu(),
opts=dict(
title='D_fake:{:d}'.format(self.global_iter)))
self.viz_train_samples.images(
self.unscale(x_real).data.cpu(),
opts=dict(
title='x_real:{:d}'.format(self.global_iter)))
self.viz_train_samples.images(
self.unscale(D_real).data.cpu(),
opts=dict(
title='D_real:{:d}'.format(self.global_iter)))
if self.visdom and self.global_iter % 10 == 0:
self.interpolation(self.fixed_z[0:1], self.fixed_z[1:2])
self.sample_img('fixed')
self.sample_img('random')
self.save_checkpoint()
if self.visdom and self.global_iter % self.timestep == 0:
# Measure of Convergence
M_global = (D_loss_real.data + abs(balance)).cpu()
X = torch.Tensor([self.global_iter])
if self.win_moc is None:
self.win_moc = self.viz_train_curves.line(
X=X,
Y=M_global,
opts=dict(title='MOC',
fillarea=True,
xlabel='iteration',
ylabel='Measure of Convergence'))
else:
self.win_moc = self.viz_train_curves.line(
X=X, Y=M_global, win=self.win_moc, update='append')
if self.global_iter % 1000 == 0:
print()
print('iter:{:d}, M:{:.3f}'.format(self.global_iter,
M_global[0]))
print(
'D_loss_real:{:.3f}, D_loss_fake:{:.3f}, G_loss:{:.3f}'
.format(D_loss_real.data[0], D_loss_fake.data[0],
G_loss.data[0]))
if self.global_iter % self.lr_step_size == 0:
self.scheduler_step()
e_elapsed = (time.time() - e_elapsed)
print()
print('epoch {:d}, [{:.2f}s]'.format(self.global_epoch, e_elapsed))
print("[*] Training Finished!")
def interpolation(self, z1, z2, n_step=10):
self.set_mode('eval')
filename = self.output_dir.joinpath('interpolation' + ':' +
str(self.global_iter) + '.jpg')
step_size = (z2 - z1) / (n_step + 1)
buff = z1
for i in range(1, n_step + 1):
_next = z1 + step_size * (i)
buff = torch.cat([buff, _next], dim=0)
buff = torch.cat([buff, z2], dim=0)
samples = self.unscale(self.G(buff))
grid = make_grid(samples.data.cpu(),
nrow=n_step + 2,
padding=1,
pad_value=0,
normalize=False)
save_image(grid, filename=filename)
if self.visdom:
self.viz_interpolations.image(grid,
opts=dict(title=str(filename),
factor=2))
self.set_mode('train')
def random_interpolation(self, n_step=10):
self.set_mode('eval')
z1 = self.sample_z(1)
z1 = Variable(cuda(z1, self.cuda))
z2 = self.sample_z(1)
z2 = Variable(cuda(z2, self.cuda))
self.interpolation(z1, z2, n_step)
self.set_mode('train')
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend="nccl",
init_method="env://")
synchronize()
args.latent = 512
args.n_mlp = 8
args.start_iter = 0
generator = Generator(
args.size,
args.latent,
args.n_mlp,
channel_multiplier=args.channel_multiplier).to(device)
discriminator = Discriminator(
args.size, channel_multiplier=args.channel_multiplier).to(device)
g_ema = Generator(args.size,
args.latent,
args.n_mlp,
channel_multiplier=args.channel_multiplier).to(device)
g_ema.eval()
accumulate(g_ema, generator, 0)
g_reg_ratio = args.g_reg_every / (args.g_reg_every + 1)
d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1)
g_optim = optim.Adam(
generator.parameters(),
lr=args.lr * g_reg_ratio,
betas=(0**g_reg_ratio, 0.99**g_reg_ratio),
)
#!./env python
from models.model import Discriminator, exhaustiveSearch
from tools.image_process import stack_svgs
from tools.common import getMaterial
import argparse
parser = argparse.ArgumentParser()
parser.add_argument(dest='query',
metavar='Query',
type=str,
help='A query sentence')
parser.add_argument('--lambda',
dest='lamb',
default=0.8,
type=float,
help='the weight between discriminators (default: 0.8)')
args = parser.parse_args()
print(args.lamb, type(args.lamb))
discriminator = Discriminator.unpickle(model_path='results/discriminator1.pkl')
layers, prob = exhaustiveSearch(args.query, discriminator, lamb=args.lamb)
print('layers: ', layers)
print('prob: %.6f' % prob, end='\n\n')
stack_svgs([getMaterial(layer) for layer in layers])
print('\n output to stack.svg')
def get_discriminator(config):
return Discriminator(config)
# dataset and dataloader
path = "D:\\python\\worksapce\\P_github_project_INFOgan_mnist_ten\\data"
dataset = MyDataset(path, opt)
dataloader = DataLoader(dataset, batch_size=opt.batchSize, shuffle=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
ngpu = int(opt.ngpu)
nz = int(opt.nz)
ngf = int(opt.ngf)
ndf = int(opt.ndf)
nc = 1
netG = Generator(ngpu, nz, ngf, nc).to(device)
netD = Discriminator(ngpu, nc, ndf).to(device)
netQ = QNet(ngpu, nc, ndf).to(device)
criterion = nn.BCELoss()
qnetloss = QNetLoss()
optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
optimizerQ = optim.Adam(netQ.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
if __name__ == "__main__":
netG.apply(weights_init)
netD.apply(weights_init)
fixed_z = Variable(torch.randn(opt.batchSize, nz - 10, 1, 1)).to(device)
fixed_c = sample_c(opt.batchSize).to(device)
# Use sigmoid activation for the last layer?
sigmoid_at_end = FG.gan in ['lsgan', 'gan']
if hasattr(FG, 'no_tanh'):
tanh_at_end = False
else:
tanh_at_end = True
G = Generator(num_channels=3,
latent_size=latent_size,
resolution=FG.target_resol,
fmap_max=latent_size,
fmap_base=8192,
tanh_at_end=tanh_at_end).to(device)
D = Discriminator(num_channels=3,
mbstat_avg=FG.mbstat_avg,
resolution=FG.target_resol,
fmap_max=latent_size,
fmap_base=8192,
sigmoid_at_end=sigmoid_at_end).to(device)
print(G)
print(D)
# exit()
if len(FG.gpu) != 1:
G = torch.nn.DataParallel(G, FG.gpu)
D = torch.nn.DataParallel(D, FG.gpu)
data = pairCelebA(FG.celeba_dir, FG.img_dir, vissum)
noise = RandomNoiseGenerator(latent_size, 'gaussian') #(32, 512)
pggan = PGGAN(G, D, data, noise, opts, vissum)
pggan.train()
def test_shape_d(self):
disc = Discriminator(self.in_channels, 8)
self.assertEqual(disc(self.x).shape, (self.N, 1, 1, 1))
print(f"{disc(self.x).shape} == {(self.N, 1, 1, 1)}")
shuffle=True)
# Test data
test_data = DatasetFromFolder(data_dir,
val_file,
params.img_types,
transform=transform)
test_data_loader = torch.utils.data.DataLoader(dataset=test_data,
batch_size=params.batch_size,
shuffle=False)
# test_input, test_target = test_data_loader.__iter__().__next__()
# Models
torch.cuda.set_device(params.gpu)
G = Generator(3, params.ngf, 3)
D = Discriminator(6, params.ndf, 1)
G.cuda()
D.cuda()
G.normal_weight_init(mean=0.0, std=0.02)
D.normal_weight_init(mean=0.0, std=0.02)
slim_params, insnorm_params = [], []
for name, param in G.named_parameters():
if param.requires_grad and name.endswith(
'weight') and 'insnorm_conv' in name:
insnorm_params.append(param)
if len(slim_params) % 2 == 0:
slim_params.append(param[:len(param) // 2])
else:
slim_params.append(param[len(param) // 2:])
fid_model = None
# defining dataset and dataloader
dataset = ImgFolderDataset(args.data_path)
dataloader = DataLoader(dataset,
batch_size=cfg.BATCH_SIZE,
shuffle=True,
num_workers=2,
drop_last=True,
pin_memory=True)
# defining models
gen = Generator(img_size=cfg.IMG_SIZE,
in_channels=cfg.IN_CHANNELS,
img_channels=cfg.CHANNELS_IMG,
z_dim=cfg.Z_DIMENSION).to(device)
dis = Discriminator(img_size=cfg.IMG_SIZE,
img_channels=cfg.CHANNELS_IMG).to(device)
# defining optimizers
opt_gen = optim.Adam(params=gen.parameters(),
lr=cfg.LEARNING_RATE,
betas=(0.0, 0.99))
opt_dis = optim.Adam(params=dis.parameters(),
lr=cfg.LEARNING_RATE,
betas=(0.0, 0.99))
# defining gradient scalers for automatic mixed precision
scaler_gen = torch.cuda.amp.GradScaler()
scaler_dis = torch.cuda.amp.GradScaler()
if args.checkpoint:
fixed_noise, cfg.START_EPOCH = load_checkpoint(args.checkpoint, gen,
opt_gen, scaler_gen,
dis, opt_dis,