def generate(dataset): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') sample_dir = os.path.join('samples', dataset) weights_dir = os.path.join('weights', dataset) os.makedirs(sample_dir, exist_ok=True) G = models.Generator(params['nz'], params['ngf'], params['image_size'], params['labels']) G.load_state_dict(torch.load(os.path.join(weights_dir, 'G.ckpt'))) G.eval() G = G.to(device) with torch.no_grad(): random_labels = torch.LongTensor( np.random.randint(0, 10, params['batch_size'])).to(device) z = torch.randn(params['batch_size'], params['nz']).to(device) fake_images = G(z, random_labels) dt_now = datetime.datetime.now() now_str = dt_now.strftime('%y%m%d%H%M%S') fake_images = fake_images.reshape(-1, 1, 28, 28) save_image(utils.denorm(fake_images), os.path.join(sample_dir, 'fake_images_{}.png'.format(now_str))) print('Saved Image ' + os.path.join(sample_dir, 'fake_images_{}.png'.format(now_str)))
def generate(): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') sample_dir = 'samples' os.makedirs(sample_dir, exist_ok=True) G = models.Generator(params['image_size'], params['latent_size'], params['hidden_size']) G.load_state_dict(torch.load('weights/G.ckpt')) G.eval() G = G.to(device) with torch.no_grad(): z = torch.randn(params['batch_size'], params['latent_size']).to(device) fake_images = G(z) fake_images = fake_images.reshape(params['batch_size'], 1, 28, 28) dt_now = datetime.datetime.now() now_str = dt_now.strftime('%y%m%d%H%M%S') save_image(utils.denorm(fake_images), os.path.join(sample_dir, 'fake_images_{}.png'.format(now_str))) print('Saved Image ' + os.path.join(sample_dir, 'fake_images_{}.png'.format(now_str)))
def train(): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') sample_dir = 'samples' os.makedirs(sample_dir, exist_ok=True) transform = transforms.Compose( [transforms.ToTensor(), transforms.Normalize(mean=[0.5], std=[0.5])]) mnist = torchvision.datasets.MNIST(root='data', train=True, transform=transform, download=True) data_loader = torch.utils.data.DataLoader(dataset=mnist, batch_size=params['batch_size'], shuffle=True) D = models.Discriminator(params['image_size'], params['hidden_size']) G = models.Generator(params['image_size'], params['latent_size'], params['hidden_size']) D = D.to(device) G = G.to(device) criterion = nn.BCELoss() d_optimizer = utils.get_optim(params, D) g_optimizer = utils.get_optim(params, G) d_losses = [] g_losses = [] total_step = len(data_loader) for epoch in range(params['epochs']): for i, (images, _) in enumerate(data_loader): # labelは使わない # (batch_size, 1, 28, 28) -> (batch_size, 1*28*28) b_size = images.size(0) images = images.reshape(b_size, -1).to(device) real_labels = torch.ones(b_size, 1).to(device) fake_labels = torch.zeros(b_size, 1).to(device) # Train discriminator outputs = D(images) d_loss_real = criterion(outputs, real_labels) real_score = outputs z = torch.randn(b_size, params['latent_size']).to(device) fake_images = G(z.detach()) outputs = D(fake_images) d_loss_fake = criterion(outputs, fake_labels) fake_score = outputs d_loss = d_loss_real + d_loss_fake d_optimizer.zero_grad() d_loss.backward() d_optimizer.step() # Train generator fake_images = G(z) outputs = D(fake_images) g_loss = criterion(outputs, real_labels) g_optimizer.zero_grad() g_loss.backward() g_optimizer.step() print( 'Epoch [{}/{}], step [{}/{}], d_loss: {:.4f}, g_loss: {:.4f}, D(x): {:.2f}, D(G(z)): {:.2f}' .format( epoch, params['epochs'], i + 1, total_step, d_loss.item(), g_loss.item(), real_score.mean().item(), fake_score.mean().item())) # .item():ゼロ次元Tensorから値を取り出す d_losses.append(d_loss.item()) g_losses.append(g_loss.item()) if (epoch + 1) == 1: images = images.reshape(b_size, 1, 28, 28) save_image(utils.denorm(images), os.path.join(sample_dir, 'real_images.png')) fake_images = fake_images.reshape(b_size, 1, 28, 28) save_image( utils.denorm(fake_images), os.path.join(sample_dir, 'fake_images-{}.png'.format(epoch + 1))) torch.save(G.state_dict(), 'weights/G.ckpt') torch.save(D.state_dict(), 'weights/D.ckpt') plt.figure(figsize=(10, 5)) plt.title("Generator and Discriminator Loss During Training") plt.plot(g_losses, label="Generator") plt.plot(d_losses, label="Discriminator") plt.xlabel("iterations") plt.ylabel("Loss") plt.legend() plt.savefig(os.path.join(sample_dir, 'loss.png'))
def train(dataset, data): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') sample_dir = os.path.join('samples', dataset) weights_dir = os.path.join('weights', dataset) os.makedirs(sample_dir, exist_ok=True) os.makedirs(weights_dir, exist_ok=True) if dataset == 'mnist': dataset = torchvision.datasets.MNIST( root=data, download=True, train=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize( # [0,1] -> [-1, 1] (0.5, ), (0.5, )), ])) params['nc'] = 1 data_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=params['batch_size'], shuffle=True) D = models.Discriminator(params['ndf'], params['image_size'], params['labels']) G = models.Generator(params['nz'], params['ngf'], params['image_size'], params['labels']) D = D.to(device) G = G.to(device) criterion = nn.BCELoss() d_optimizer = utils.get_optim(params, D) g_optimizer = utils.get_optim(params, G) d_losses = [] g_losses = [] total_step = len(data_loader) for epoch in range(params['epochs']): for i, (images, labels) in enumerate(data_loader): b_size = images.size(0) images = images.reshape(b_size, -1).to(device) labels = labels.to(device) real_labels = torch.ones(b_size).to(device) fake_labels = torch.zeros(b_size).to(device) random_labels = torch.LongTensor(np.random.randint( 0, 10, b_size)).to(device) # Train discriminator outputs = D(images, labels) d_loss_real = criterion(outputs, real_labels) real_score = outputs z = torch.randn(b_size, params['nz']).to(device) fake_images = G(z, random_labels) outputs = D(fake_images.detach(), random_labels) d_loss_fake = criterion(outputs, fake_labels) fake_score = outputs d_loss = d_loss_real + d_loss_fake d_optimizer.zero_grad() d_loss.backward() d_optimizer.step() # Train generator # ランダムなラベルを再定義 random_labels = torch.LongTensor(np.random.randint( 0, 10, b_size)).to(device) fake_images = G(z, random_labels) outputs = D(fake_images, random_labels) g_loss = criterion(outputs, real_labels) g_optimizer.zero_grad() g_loss.backward() g_optimizer.step() print( 'Epoch [{}/{}], step [{}/{}], d_loss: {:.4f}, g_loss: {:.4f}, D(x): {:.2f}, D(G(z)): {:.2f}' .format( epoch, params['epochs'], i + 1, total_step, d_loss.item(), g_loss.item(), real_score.mean().item(), fake_score.mean().item())) # .item():ゼロ次元Tensorから値を取り出す g_losses.append(g_loss.item()) d_losses.append(d_loss.item()) if (epoch + 1) == 1: images = images.reshape(b_size, 1, 28, 28) save_image(utils.denorm(images), os.path.join(sample_dir, 'real_images.png')) fake_images = fake_images.reshape(b_size, 1, 28, 28) save_image( utils.denorm(fake_images), os.path.join(sample_dir, 'fake_images-{}.png'.format(epoch + 1))) torch.save(G.state_dict(), os.path.join(weights_dir, 'G.ckpt')) torch.save(D.state_dict(), os.path.join(weights_dir, 'D.ckpt')) plt.figure(figsize=(10, 5)) plt.title("Generator and Discriminator Loss During Training") plt.plot(g_losses, label="Generator") plt.plot(d_losses, label="Discriminator") plt.xlabel("iterations") plt.ylabel("Loss") plt.legend() plt.savefig(os.path.join(sample_dir, 'loss.png'))