def train(args):
declare_global_parameter(args)
netG = Generator(NZ + NO * 5, NC, DIM, IMAGE_SIZE, 0)
netD = Discriminator(NZ, NC, DIM, IMAGE_SIZE, 0)
netC = Critic(NZ, NC, DIM, IMAGE_SIZE, 0)
netE = models.encoder.Encoder(NZ, NC, NO, DIM, IMAGE_SIZE // 2, 0)
print(netG)
print(netC)
print(netD)
print(netE)
if CUDA:
netG.cuda()
netC.cuda()
netD.cuda()
netE.cuda()
if NGPU > 1:
netG = nn.DataParallel(netG, device_ids=range(NGPU))
netC = nn.DataParallel(netC, device_ids=range(NGPU))
netD = nn.DataParallel(netD, device_ids=range(NGPU))
netE = nn.DataParallel(netE, device_ids=range(NGPU))
cudnn.benchmark = True
optimizerD = optim.Adam(netD.parameters(), lr=LR_D, betas=(0.5, 0.9))
optimizerG = optim.Adam(netG.parameters(), lr=LR_G, betas=(0.5, 0.9))
optimizerC = optim.Adam(netC.parameters(), lr=LR_D, betas=(0.5, 0.9))
optimizerE = optim.Adam(netE.parameters(), lr=LR_G, betas=(0.5, 0.9))
# Dataset loader
transform = tv.transforms.Compose(
[tv.transforms.Scale(IMAGE_SIZE),
tv.transforms.ToTensor()])
if DATASET == 'mnist':
dataset = tv.datasets.MNIST(DATA_SAVE,
train=True,
transform=transform,
download=True)
if CUDA:
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=BATCH_SIZE * 2,
shuffle=True,
num_workers=NW,
pin_memory=True,
drop_last=True)
else:
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=BATCH_SIZE * 2,
shuffle=True,
num_workers=NW,
drop_last=True)
for epoch in range(EPOCHS):
run_epoch(dataloader, netG, netE, netC, netD, optimizerG, optimizerE,
optimizerC, optimizerD)
def main():
# init random seed
init_random_seed(params.manual_seed)
#check the needed dirs of config
check_dirs()
cudnn.benchmark = True
torch.cuda.set_device(params.gpu_id[0]) #set current device
print('=== Build model ===')
#gpu mode
generator = Generator()
discriminator = Discriminator()
generator = nn.DataParallel(generator, device_ids=params.gpu_id).cuda()
discriminator = nn.DataParallel(discriminator,
device_ids=params.gpu_id).cuda()
# restore trained model
if params.generator_restored:
generator = restore_model(generator, params.generator_restored)
if params.discriminator_restored:
discriminator = restore_model(discriminator,
params.discriminator_restored)
# container of training
trainer = Trainer(generator, discriminator)
if params.mode == 'train':
# data loader
print('=== Load data ===')
train_dataloader = get_data_loader(params.dataset)
print('=== Begin training ===')
trainer.train(train_dataloader)
print('=== Generate {} images, saving in {} ==='.format(
params.num_images, params.save_root))
trainer.generate_images(params.num_images, params.save_root)
elif params.mode == 'test':
if params.generator_restored:
print('=== Generate {} images, saving in {} ==='.format(
params.num_images, params.save_root))
trainer.generate_images(params.num_images, params.save_root)
else:
assert False, '[*]load Generator model first!'
else:
assert False, "[*]mode must be 'train' or 'test'!"
def get_gan(gan_type: GANType,
device: torch.device,
n_power_iterations: int = None) -> Tuple[nn.Module, nn.Module]:
r"""Fetching GAN and moving it to proper device.
Args:
-gan_type (GANType): DCGAN or SN-DCGAN.
-device (torch.device): On which device (eg. GPU) to move models.
-n_power_iterations (int): Number of iterations for l_2 matrix norm.
Returns:
-G (nn.Module): Generator.
-D (nn.Module): Discriminator.
"""
if gan_type == GANType.DCGAN:
G = Generator().to(device)
D = Discriminator().to(device)
elif gan_type == GANType.SN_DCGAN:
G = SNGenerator().to(device)
D = SNDiscriminator(n_power_iterations).to(device)
return G, D
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from models.dcgan import Generator, Discriminator, weights_init
from train_gan import Trainer
if __name__ == '__main__':
p2data = "/hdd1/diploma_outputs/outputs"
dset_name = "train_flowers.hdf5"
imageSize = 64
transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomResizedCrop(size=imageSize, scale=(0.9, 0.95), ratio=(1, 1)),
transforms.ToTensor()])
batchSize = 256
workers = 4
dset = Text2ImageDataset(os.path.join(p2data, dset_name), split="train", transform=transform, mean=0, std=1)
gen = Generator()
discr = Discriminator()
# gen, discr, type, dataset, lr, diter, vis_screen, save_path, l1_coef, l2_coef,
# pre_trained_gen,
# pre_trained_disc, batch_size, num_workers, epochs
gan_trainer = Trainer(gen, discr, dset, 0.0002, "gan4", "output1", 50, 100, False, False, 64, 4, 30)
gan_trainer.train(cls=True, spe=125)
# custom weights initialization called on netG and netD
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
m.weight.data.normal_(0.0, 0.02)
elif classname.find('BatchNorm') != -1:
m.weight.data.normal_(1.0, 0.02)
m.bias.data.fill_(0)
netG = Generator().to(device)
netG.apply(weights_init)
if opt.netG != '':
netG.load_state_dict(torch.load(opt.netG))
print(netG)
netD = Discriminator().to(device)
netD.apply(weights_init)
if opt.netD != '':
netD.load_state_dict(torch.load(opt.netD))
print(netD)
criterion = nn.BCELoss()
fixed_noise = torch.randn(opt.batchSize, nz, 1, 1, device=device)
real_label = 1
fake_label = 0
# setup optimizer
optimizerD = optim.Adam(netD.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
epochs = opt.epochs
save_epoch_interval = opt.save_epoch_interval
train_D_iter = opt.train_D_iter
# persistence related parameters
writer = SummaryWriter(save_path)
utiler = TensorImageUtils(save_path)
nrow = opt.nrow
# dataset
data = choose_dataset(opt)
in_channels = opt.in_channels
# models
dim_z = opt.dim_z
netD = Discriminator(in_channels, dim_z)
netG = Generator(in_channels, dim_z)
headD = DHead(512)
headQ = QHead(512, num_clsses, num_continuous_variables)
if use_cuda:
netD.cuda()
netG.cuda()
headD.cuda()
headQ.cuda()
# training config
optimizer_D = optim.Adam([{
"params": netD.parameters()
}, {
"params": headD.parameters()
def main():
load_pretrained = False
if os.path.isfile(os.path.join(config['pretrained'] + '_netG.pt')):
load_pretrained = True
netD_path = os.path.join(config['pretrained'] + '_netD.pt')
netG_path = os.path.join(config['pretrained'] + '_netG.pt')
current_epoch = int(config['pretrained'].split(os.path.sep)[-1].split("_")[0]) + 1
current_iter = int(config['pretrained'].split(os.path.sep)[-1].split("_")[1])
print(current_epoch, current_iter)
print("pretrained")
else:
current_epoch = 0
dataset = PokemonDataset(dataroot=config['dataroot'],
transform=transforms.Compose([
transforms.Resize(int(config['image_size'])),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0),
transforms.RandomRotation(10),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]),
config=config)
# Create the dataloader
dataloader = torch.utils.data.DataLoader(dataset, batch_size= int(config['batch_size']),
shuffle=True, num_workers= int(config['workers']))
device = torch.device("cuda:0" if (torch.cuda.is_available() and int(config['ngpu']) > 0) else "cpu")
# Create the generator
netG = Generator(config).to(device)
# Handle multi-gpu if desired
if (device.type == 'cuda') and (int(config['ngpu']) > 1):
netG = nn.DataParallel(netG, list(range(int(config['ngpu']))))
# Apply the weights_init function to randomly initialize all weights
# to mean=0, stdev=0.2.
if load_pretrained:
netG.load_state_dict(torch.load(netG_path))
else:
netG.apply(weights_init)
netG.train()
# Print the model
print(netG)
# Create the discriminator
netD = Discriminator(config).to(device)
# Handle multi-gpu if desired
if (device.type == 'cuda') and (int(config['ngpu']) > 1):
netD = nn.DataParallel(netD, list(range(int(config['ngpu']))))
# Apply the weights_init function to randomly initialize all weights
# to mean=0, stdev=0.2.
# netD.apply(weights_init)
if load_pretrained:
netD.load_state_dict(torch.load(netD_path))
else:
netD.apply(weights_init)
netD.train()
# Print the model
print(netD)
# Initialize BCELoss function
criterion = nn.BCELoss()
# Create batch of latent vectors that we will use to visualize
# the progression of the generator
fixed_noise = torch.randn(64, int(config['nz']), 1, 1, device=device)
# Establish convention for real and fake labels during training
real_label = 0.9 # GAN tricks #1: label smoothing
fake_label = 0
# Setup Adam optimizers for both G and D
# optimizerD = optim.Adam(netD.parameters(), lr=float(config['netD_lr']), betas=(float(config['beta1']), 0.999))
optimizerD = optim.Adam(filter(lambda p: p.requires_grad, netD.parameters()), lr=float(config['netD_lr']), betas=(float(config['beta1']), 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=float(config['netG_lr']), betas=(float(config['beta1']), 0.999))
# Training Loop
num_epochs = int(config['num_epochs'])
nz = int(config['nz'])
# Lists to keep track of progress
img_list = []
G_losses = []
D_losses = []
frames = []
iters = 0
if load_pretrained:
iters = current_iter
print("Starting Training Loop...")
start_time = time.time()
# For each epoch
for epoch in range(current_epoch, num_epochs):
# For each batch in the dataloader
for i, data in enumerate(dataloader, 0):
############################
# (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
###########################
## Train with all-real batch
netD.zero_grad()
# Format batch
real_cpu = data[0].to(device)
b_size = real_cpu.size(0)
label = torch.full((b_size,), real_label, device=device)
# Forward pass real batch through D
output = netD(real_cpu).view(-1)
# Calculate loss on all-real batch
errD_real = criterion(output, label)
# Calculate gradients for D in backward pass
errD_real.backward()
D_x = output.mean().item()
## Train with all-fake batch
# Generate batch of latent vectors
noise = torch.randn(b_size, nz, 1, 1, device=device)
# Generate fake image batch with G
fake = netG(noise)
label.fill_(fake_label)
# Classify all fake batch with D
output = netD(fake.detach()).view(-1)
# Calculate D's loss on the all-fake batch
errD_fake = criterion(output, label)
# Calculate the gradients for this batch
errD_fake.backward()
D_G_z1 = output.mean().item()
# Add the gradients from the all-real and all-fake batches
errD = errD_real + errD_fake
# Update D
optimizerD.step()
############################
# (2) Update G network: maximize log(D(G(z)))
###########################
netG.zero_grad()
label.fill_(real_label) # fake labels are real for generator cost
# Since we just updated D, perform another forward pass of all-fake batch through D
output = netD(fake).view(-1)
# Calculate G's loss based on this output
errG = criterion(output, label)
# Calculate gradients for G
errG.backward()
D_G_z2 = output.mean().item()
# Update G
optimizerG.step()
# Output training stats
if i % 50 == 0:
end_time = time.time()
duration = end_time - start_time
print(f"{duration:.2f}s, [{epoch}/{num_epochs}][{i}/{len(dataloader)}]\tLoss_D: {errD.item():.4f}\t \
Loss_G: {errG.item():.4f}\tD(x): {D_x:.4f}\tD(G(z)): {D_G_z1:.4f} / {D_G_z2:.4f}")
start_time = time.time()
# Save Losses for plotting later
G_losses.append(errG.item())
D_losses.append(errD.item())
# Check how the generator is doing by saving G's output on fixed_noise
if (iters % int(config['save_freq']) == 0) or ((epoch == num_epochs-1) and (i == len(dataloader)-1)):
with torch.no_grad():
fake = netG(fixed_noise).detach().cpu()
grid = vutils.make_grid(fake, padding=2, normalize=True)
ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
im = Image.fromarray(ndarr)
im.save(os.path.join("output", f"epoch{epoch}_iter{iters}.png"))
frames.append(im)
torch.save(netD.state_dict(), os.path.join("output", f"{epoch}_{iters}_netD.pt"))
torch.save(netG.state_dict(), os.path.join("output", f"{epoch}_{iters}_netG.pt"))
iters += 1
plt.figure(figsize=(10,5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(G_losses,label="G")
plt.plot(D_losses,label="D")
plt.xlabel("iterations")
plt.ylabel("Loss")
plt.legend()
plt.savefig(os.path.join("output", "loss_curve.png"))
frames[0].save(os.path.join('output', 'animation.gif'), format='GIF', append_images=frames[1:], save_all=True, duration=500, loop=0)
# Training settings
parser = argparse.ArgumentParser(description='PyTorch Refine GAN')
parser.add_argument('--cuda', action='store_true', default=True)
args = parser.parse_args()
dataset_train = MaterialsDataset("PBR_dataset_256/", test = True)
dataset_test = MaterialsDataset("PBR_dataset_256/", test = False)
train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=16, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=16, shuffle=True)
if not torch.cuda.is_available():
print("Warning cuda not found, falling back on CPU!")
args.cuda = False
G = UNet(n_channels=3, n_output=3)
D = Discriminator()
if args.cuda:
G.cuda()
D.cuda()
lr=0.001
# run forever
for epoch in range(10000):
lr *= 0.9995
print("Epoch {}, lr: {}".format(epoch,lr) )
G_solver = torch.optim.Adam(model.parameters(),lr=lr )
D_solver = torch.optim.Adam(model.parameters(),lr=lr )
train_loss = list()
def main():
params = parseyaml()
if params['arch'] == 'Generator':
device = to_gpu(ngpu=params['n_gpu'])
if params['image_size'] == 64:
netG = Generator(ngpu=0, nz=256,
ngf=64, nc=64).to(device)
elif params['image_size'] == 128:
netG = Generator_128(ngpu=0, nz=256,
ngf=64, nc=64).to(device)
elif params['image_size'] == 256:
netG = Generator_256(ngpu=0, nz=256,
ngf=64, nc=64).to(device)
netG.apply(weights_init)
netG.load_state_dict(torch.load(params['path']))
for i in range(params['quantity']):
fixed_noise = torch.randn(64, 256, 1, 1, device=device)
fakes = netG(fixed_noise)
for j in range(len(fakes)):
save_image(fakes[j], params['out'] + params['run'] +
'_' + str(i) + '_' + str(j) + '_img.png')
else:
dataloader = dataLoader(
path=params['path'], image_size=params['image_size'], batch_size=params['batch_size'],
workers=params['loader_workers'])
device = to_gpu(ngpu=params['n_gpu'])
if params['arch'] == 'DCGAN':
if params['image_size'] == 64:
netG = Generator(ngpu=params['n_gpu'], nz=params['latent_vector'],
ngf=params['gen_feature_maps'], nc=params['number_channels']).to(device)
netD = Discriminator(params['n_gpu'], nc=params['number_channels'],
ndf=params['dis_feature_maps']).to(device)
elif params['image_size'] == 128:
netG = Generator_128(ngpu=params['n_gpu'], nz=params['latent_vector'],
ngf=params['gen_feature_maps'], nc=params['number_channels']).to(device)
netD = Discriminator_128(params['n_gpu'], nc=params['number_channels'],
ndf=params['dis_feature_maps']).to(device)
elif params['image_size'] == 256:
netG = Generator_256(ngpu=params['n_gpu'], nz=params['latent_vector'],
ngf=params['gen_feature_maps'], nc=params['number_channels']).to(device)
netD = Discriminator_256(params['n_gpu'], nc=params['number_channels'],
ndf=params['dis_feature_maps']).to(device)
elif params['arch'] == 'SNGAN':
if params['image_size'] == 64:
netG = Generator(ngpu=params['n_gpu'], nz=params['latent_vector'],
ngf=params['gen_feature_maps'], nc=params['number_channels']).to(device)
netD = Discriminator_SN(params['n_gpu'], nc=params['number_channels'],
ndf=params['dis_feature_maps']).to(device)
elif params['image_size'] == 128:
netG = Generator_128(ngpu=params['n_gpu'], nz=params['latent_vector'],
ngf=params['gen_feature_maps'], nc=params['number_channels']).to(device)
netD = Discriminator_SN_128(params['n_gpu'], nc=params['number_channels'],
ndf=params['dis_feature_maps']).to(device)
elif params['image_size'] == 256:
netG = Generator_256(ngpu=params['n_gpu'], nz=params['latent_vector'],
ngf=params['gen_feature_maps'], nc=params['number_channels']).to(device)
netD = Discriminator_SN_256(params['n_gpu'], nc=params['number_channels'],
ndf=params['dis_feature_maps']).to(device)
if (device.type == 'cuda') and (params['n_gpu'] > 1):
netG = nn.DataParallel(netG, list(range(params['n_gpu'])))
if (device.type == 'cuda') and (params['n_gpu'] > 1):
netD = nn.DataParallel(netD, list(range(params['n_gpu'])))
netG.apply(weights_init)
netD.apply(weights_init)
print(netG)
print(netD)
criterion = nn.BCELoss()
fixed_noise = torch.randn(params['image_size'],
params['latent_vector'], 1, 1, device=device)
if params['learning_rate'] >= 1:
optimizerD = optim.Adam(netD.parameters(), lr=0.0002 * params['learning_rate'], betas=(
params['beta_adam'], 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=0.0002, betas=(
params['beta_adam'], 0.999))
else:
optimizerD = optim.Adam(netD.parameters(), lr=params['learning_rate'], betas=(
params['beta_adam'], 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=params['learning_rate'], betas=(
params['beta_adam'], 0.999))
G_losses, D_losses, img_list, img_list_only = training_loop(num_epochs=params['num_epochs'], dataloader=dataloader,
netG=netG, netD=netD, device=device, criterion=criterion, nz=params[
'latent_vector'],
optimizerG=optimizerG, optimizerD=optimizerD, fixed_noise=fixed_noise, out=params['out'] + params['run'] + '_')
loss_plot(G_losses=G_losses, D_losses=D_losses, out=params['out'] + params['run'] + '_')
image_grid(dataloader=dataloader, img_list=img_list,
device=device, out=params['out'] + params['run'] + '_')
compute_metrics(real=next(iter(dataloader)), fakes=img_list_only,
size=params['image_size'], out=params['out'] + params['run'] + '_')
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1: # Conv weight init
m.weight.data.normal_(0.0, 0.02)
elif classname.find('BatchNorm') != -1: # BatchNorm weight init
m.weight.data.normal_(1.0, 0.02)
m.bias.data.fill_(0)
G = Generator(latent_dim, ngf, channels).cuda(device)
D = Discriminator(ndf, channels).cuda(device)
G.apply(weights_init)
D.apply(weights_init)
dataloader = torch.utils.data.DataLoader(
dataset=torchvision.datasets.ImageFolder(root='C:/Users/우리집/MyWGAN/images',
transform=trans),
batch_size=batch_size)
Optimizer_D = torch.optim.RMSprop(D.parameters(), lr=lr)
Optimizer_G = torch.optim.RMSprop(G.parameters(), lr=lr)
# Optimizer_G = torch.optim.Adam(G.parameters(),lr=lr, betas=(beta1, 0.999))
# Optimizer_D = torch.optim.Adam(D.parameters(),lr=lr, betas=(beta1, 0.999))
G_losses = []