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,