def __init__(
self,
generator_and_opt: [Generator, torch.optim] = None,
discriminator_and_opt: [Discriminator, torch.optim] = None,
input_size: int = None,
hidden_channel: int = 128,
latent_dim: int = 100,
learning_rate: float = 1e-4,
):
self.generator = None
self.discriminator = None
super().__init__()
# Generator
if generator_and_opt is None:
assert input_size is None, "generator_and_opt or input_size should be given."
if self.generator is not None:
self.generator = Generator(input_size=input_size,
latent_dim=latent_dim,
hidden_channel=hidden_channel)
self.generator_opt = torch.optim.Adam(
self.generator.parameters(), learning_rate)
else:
self.generator, self.generator_opt = generator_and_opt
# Discriminator
if discriminator_and_opt is None:
assert input_size is None, "discriminator_and_opt or input_size should be given."
if self.discriminator is not None:
self.discriminator = Discriminator(
input_size=input_size, hidden_channel=hidden_channel)
self.discriminator_opt = torch.optim.Adam(
self.discriminator.parameters(), learning_rate)
else:
self.discriminator, self.discriminator_opt = discriminator_and_opt
def __init__(self, args):
self.args = args
self.pretrained = False
self.epoch = 0
self.G = Generator()
self.D = Discriminator()
self.g_optimizer = optim.Adam(self.G.parameters(), lr=1E-4)
self.d_optimizer = optim.Adam(self.D.parameters(), lr=1E-4)
self.g_scheduler = optim.lr_scheduler.StepLR(self.g_optimizer,
step_size=40)
self.d_scheduler = optim.lr_scheduler.StepLR(self.d_optimizer,
step_size=40)
self.train_losses = []
self.val_losses = []
if args.load_model:
self._load_state(args.load_model)
# extract all layers prior to the last softmax of VGG-19
vgg19_layers = list(models.vgg19(pretrained=True).features)[:36]
self.vgg19 = nn.Sequential(*vgg19_layers).eval()
for param in self.vgg19.parameters():
param.requires_grad = False
self.mse_loss = torch.nn.MSELoss()
self.bce_loss = torch.nn.BCELoss()
def build_model(self):
self.G = Generator()
self.D = Discriminator()
self.G_optim = optim.Adam(self.G.parameters(), self.G_lr,
(self.beta1, 0.999))
self.D_optim = optim.Adam(self.D.parameters(), self.D_lr,
(self.beta1, 0.999))
if self.loss_type == 'BCEwL':
self.criterion = nn.BCEWithLogitsLoss()
elif self.loss_type == 'WGAN':
pass
elif self.loss_type == 'WGAN+':
pass
self.fixed_sample = None
self.fixed_noise = None
# self.true = torch.ones([self.batch_size, 1, 1, 1], requires_grad=False).to(self.device)
# self.false = torch.zeros([self.batch_size, 1, 1, 1], requires_grad=False).to(self.device)
# Change to GPU mode
print('Change CPU mode to GPU mode...')
self.G.to(self.device)
self.D.to(self.device)
print('Creating models are success...')
def __init__(self, config):
super(Transformer_EncoderDecoder, self).__init__()
c = copy.deepcopy
self.attn = MultiHeadedAttention(config['head'], config['emb_dim'])
self.ff = PositionwiseFeedForward(config['emb_dim'], config['d_ff'],
config['drop_out'])
self.position = PositionalEncoding(config['emb_dim'],
config['drop_out'])
self.encoder = Encoder(
EncoderLayer(config['emb_dim'], c(self.attn), c(self.ff),
config['drop_out']), config['N_layers'])
self.decoder = Decoder(
DecoderLayer(config['emb_dim'], c(self.attn), c(self.attn),
c(self.ff), config['drop_out']), config['N_layers'])
self.src_embed = nn.Sequential(
Embeddings(config['emb_dim'], config['vocab_size']),
c(self.position))
self.tgt_embed = nn.Sequential(
Embeddings(config['emb_dim'], config['vocab_size']),
c(self.position))
self.generator = Generator(config['emb_dim'], config['vocab_size'])
self.fc_out = nn.Linear(config['emb_dim'], config['vocab_size'])
self.model = EncoderDecoder(self.encoder, self.decoder, self.src_embed,
self.tgt_embed, self.generator)
def main():
args = parse_args()
model = Model()
with open(args.model, 'rb') as file:
model.load(file)
generator = Generator(model)
if args.output is not None:
file = open(args.output, 'a', encoding='utf-8')
else:
file = stdout
for i in range(0, args.count):
text = generator.generate(args.length,
seed=args.seed,
min_n=args.min_n,
n=args.n,
break_on_end=args.break_on_end)
if args.wrap:
text = '\n'.join(wrap(text))
print(text, '\n', file=file)
if args.output is not None:
file.close()
def build_model(self):
self.G = Generator(image_size=self.imsize,
z_dim=self.z_dim,
conv_dim=self.g_dim)
self.D = Discriminator(conv_dim=self.d_dim)
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])
self.print_network(self.G, 'G')
self.print_network(self.D, 'D')
self.G.to(self.device)
self.D.to(self.device)
def train(dataset_dir, output_dir):
"""Train discriminator and generator"""
if torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
device = torch.device('cuda')
else:
device = torch.device('cpu')
discriminator = Discriminator()
discriminator.to(device)
discriminator.train()
generator = Generator()
generator.to(device)
generator.train()
optimizer_discriminator = torch.optim.Adam(discriminator.parameters(), lr=0.0001)
optimizer_generator = torch.optim.Adam(generator.parameters(), lr=0.0001)
loss_func = nn.BCELoss()
loader = LFWLoader(dataset_dir)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
for epoch in range(10):
for image in tqdm(loader):
# train discriminator on true
outputs = discriminator(torch.tensor(image, dtype=torch.float32) / 255)
loss = loss_func(outputs, torch.tensor([1.0]))
optimizer_discriminator.zero_grad()
loss.backward()
optimizer_discriminator.step()
# train discriminator on false
outputs = discriminator(generator(generate_random_seed(100)).detach())
loss = loss_func(outputs, torch.tensor([0.0]))
optimizer_discriminator.zero_grad()
loss.backward()
optimizer_discriminator.step()
# train generator
outputs = discriminator(generator(generate_random_seed(100)))
loss = loss_func(outputs, torch.tensor([1.0]))
optimizer_generator.zero_grad()
loss.backward()
optimizer_generator.step()
torch.save(generator, os.path.join(output_dir, 'G_%d.pt' % epoch))
torch.save(discriminator, os.path.join(output_dir, 'D_%d.pt' % epoch))
def __init__(self,config):
super(BiGAN,self).__init__()
self._work_type = config.work_type
self._epochs = config.epochs
self._batch_size = config.batch_size
self._encoder_lr = config.encoder_lr
self._generator_lr = config.generator_lr
self._discriminator_lr = config.discriminator_lr
self._latent_dim = config.latent_dim
self._weight_decay = config.weight_decay
self._img_shape = (config.input_size,config.input_size)
self._img_save_path = config.image_save_path
self._model_save_path = config.model_save_path
self._device = config.device
if self._work_type == 'train':
# Loss function
self._adversarial_criterion = torch.nn.MSELoss()
# Initialize generator, encoder and discriminator
self._G = Generator(self._latent_dim,self._img_shape).to(self._device)
self._E = Encoder(self._latent_dim,self._img_shape).to(self._device)
self._D = Discriminator(self._latent_dim,self._img_shape).to(self._device)
self._G.apply(self.weights_init)
self._E.apply(self.weights_init)
self._D.apply(self.discriminator_weights_init)
self._G_optimizer = torch.optim.Adam([{'params' : self._G.parameters()},{'params' : self._E.parameters()}],
lr=self._generator_lr,betas=(0.5,0.999),weight_decay=self._weight_decay)
self._D_optimizer = torch.optim.Adam(self._D.parameters(),lr=self._discriminator_lr,betas=(0.5,0.999))
self._G_scheduler = lr_scheduler.ExponentialLR(self._G_optimizer, gamma= 0.99)
self._D_scheduler = lr_scheduler.ExponentialLR(self._D_optimizer, gamma= 0.99)
def build_model(self):
"""Build generator and discriminator."""
self.generator = Generator(z_dim=self.z_dim)
print(count_parameters(self.generator))
self.discriminator = Discriminator()
print(count_parameters(self.discriminator))
self.g_optimizer = optim.Adam(self.generator.parameters(),
self.lr, (self.beta1, self.beta2))
self.d_optimizer = optim.Adam(self.discriminator.parameters(),
self.lr*1, (self.beta1, self.beta2))
if self.epoch:
g_path = os.path.join(self.model_path, 'generator-%d.pkl' % self.epoch)
d_path = os.path.join(self.model_path, 'discriminator-%d.pkl' % self.epoch)
g_optim_path = os.path.join(self.model_path, 'gen-optim-%d.pkl' % self.epoch)
d_optim_path = os.path.join(self.model_path, 'dis-optim-%d.pkl' % self.epoch)
self.generator.load_state_dict(torch.load(g_path))
self.discriminator.load_state_dict(torch.load(d_path))
self.g_optimizer.load_state_dict(torch.load(g_optim_path))
self.d_optimizer.load_state_dict(torch.load(d_optim_path))
if torch.cuda.is_available():
self.generator.cuda()
self.discriminator.cuda()
def run(argv):
global args
global model
global loader
global test_loader
global output_shape
global visualization_mode
parser = argparse.ArgumentParser()
# training
parser.add_argument('--train', action='store_true', default=False, help='Train a model')
parser.add_argument('--train_generator', type=bool, default=True, help='Include generator training')
parser.add_argument('--save_to', type=str, default='model.bin', help='Where to save the trained model')
parser.add_argument('--batch_size', type=int, default=2048, help='Training batch size')
parser.add_argument('--epochs', type=int, default=200, help='Number of epochs to train')
parser.add_argument('--lr', type=float, default=0.001, help='Learning rate')
parser.add_argument('--lr_step', type=int, default=0, help='Learning rate step size')
parser.add_argument('--lr_step_gamma', type=float, default=0.5, help='Learning rate step gamma')
parser.add_argument('--sigma', type=float, default=0.2, help='Noise standard deviation')
parser.add_argument('--reduce_sigma_every', type=int, default=0, help='Decrease sigma value every N batches')
parser.add_argument('--reduce_sigma_gamma', type=float, default=0.8, help='When decreasing sigma, sigma *= sigma_gamma')
parser.add_argument("--min_sigma", type=float, default=0, help='Lower bound of sigma')
parser.add_argument('--stop_training_real_dde_after', type=int, default=-1, help='Only train real_dde for this number of epochs')
parser.add_argument('--train_gen_every', type=int, default=10, help='Step generator every N iterations')
parser.add_argument('--visualize_every', type=int, default=50, help='Visualize model every N batches')
# loading
parser.add_argument('--load_from', type=str, default=None, help='Where to load a trained model')
# dataset
parser.add_argument('--dataset', type=str, default='ts', help='Dataset (ts, eg, cb, mnist, fashion, stacked-mnist, ncsn)')
parser.add_argument('--dataset_dir', type=str, default='/tmp/', help='Dataset location for MNIST and Fashion. Defaults to /tmp/')
parser.add_argument('--download', action='store_true', default=False, help='Download the dataset if it does not exist')
parser.add_argument('--datasize', type=int, default=51200, help='Datasize (pseudo, for epoch computation) for 2D datasets')
parser.add_argument('--dynamic', action='store_true', default=True, help='Whether 2D datasets are dynamically generated')
# model parameters
parser.add_argument('--dense_sigmoid', action='store_true', default=False, help='Add sigmoid to DenseNet generator')
parser.add_argument('--gen_activation', type=str, default='softplus', help='Generator activation function')
parser.add_argument('--dde_activation', type=str, default='softplus', help='Discriminator activation function')
parser.add_argument('--n_input', type=int, default=2, help='Number of input to generator')
parser.add_argument('--n_hidden', type=int, default=32, help='Number of hidden units in a layer')
parser.add_argument('--n_gen_hidden', type=int, default=64, help='Number of hidden units in DenseNet generator layer')
parser.add_argument('--n_layers', type=int, default=25, help='Number of layers')
parser.add_argument('--n_feat_gen', type=int, default=64, help='Number of features in DCGAN generator')
parser.add_argument('--n_feat_dsc', type=int, default=45, help='Number of features in DCGAN discriminator')
parser.add_argument('--gen_bn', type=bool, default=True, help='Use BN in DCGAN generator')
parser.add_argument('--dsc_activation', type=str, default='lrelu:0.2', help='Activation in DCGAN discriminator')
parser.add_argument('--device', type=str, default='cuda:0', help='PyTorch device')
# visualization
parser.add_argument('--visualize', action='store_true', default=False, help='Visualize a model')
parser.add_argument('--n_toy_samples', type=int, default=50000, help='Number of samples to visualize for toy datasets')
parser.add_argument('--save_toy_samples_to', type=str, default=None, help='If specified, save generated toy samples to file')
parser.add_argument('--vis_path', type=str, default='', help='Path prefix to save visualization figures.')
parser.add_argument('--n_uci_samples', type=int, default=51200, help='Number of samples to use when computing UCI ALL')
args = parser.parse_args(argv) if not argv is None else parser.parse_args()
if args.dataset == 'ts':
loader = twoSpiralsLoader(args.datasize, args.batch_size, dynamic=args.dynamic)
args.type = 'mlp'
elif args.dataset == 'eg':
loader = eightGaussiansLoader(args.datasize, args.batch_size, dynamic=args.dynamic)
args.type = 'mlp'
elif args.dataset == 'cb':
loader = checkerboardLoader(args.datasize, args.batch_size, dynamic=args.dynamic)
args.type = 'mlp'
elif args.dataset == 'mnist':
loader = mnistLoader(args.dataset_dir, args.download, args.batch_size)
args.type = 'densenet'
visualization_mode = 'images'
elif args.dataset == 'fashion':
loader = fashionLoader(args.dataset_dir, args.download, args.batch_size)
args.type = 'densenet'
visualization_mode = 'images'
elif args.dataset == 'stacked-mnist':
loader = stackedMnistLoader(args.dataset_dir, args.download, args.batch_size)
args.type = 'dcgan'
visualization_mode = 'images'
elif args.dataset == 'ncsn':
loader = stackedMnistLoader(args.dataset_dir, args.download, args.batch_size)
args.type = 'ncsn'
visualization_mode = 'images'
elif args.dataset.startswith('uci_'):
name = args.dataset[4:]
loader, test_loader = uciLoader(args.dataset_dir, name, args.batch_size)
args.type = 'mlp'
visualization_mode = 'numbers'
else:
print(f'Unknown dataset: {args.dataset}')
return
output_shape = list(loader.dataset[0][0].shape)
print(f'Output shape: {output_shape}')
output_length = loader.dataset[0][0].view(-1).shape[0]
print(f'Output length: {output_length}')
gen_activation = activation_from_name(args.gen_activation)
dde_activation = activation_from_name(args.dde_activation)
if args.type == 'mlp':
model = DdeModel(MlpModule(output_length, args.n_hidden, args.n_layers, output_length, gen_activation, True, args.device),
MlpModule(output_length, args.n_hidden, args.n_layers, 1, dde_activation, False, args.device),
MlpModule(output_length, args.n_hidden, args.n_layers, 1, dde_activation, False, args.device))
elif args.type == 'densenet':
end_activation = nn.Sigmoid if args.dense_sigmoid else None
model = DdeModel(DenseNetModule(args.n_input, args.n_gen_hidden, args.n_layers, output_length, gen_activation, end_activation, args.device),
DenseNetModule(output_length, args.n_hidden, args.n_layers, 1, dde_activation, None, args.device),
DenseNetModule(output_length, args.n_hidden, args.n_layers, 1, dde_activation, None, args.device))
elif args.type == 'dcgan':
dsc_activation = activation_from_name(args.dsc_activation)
model = DdeModel(Generator(args.n_input, args.n_feat_gen, output_shape[0], args.gen_bn, args.device),
Discriminator(args.n_feat_dsc, output_shape[0], dsc_activation, None, False, args.device),
Discriminator(args.n_feat_dsc, output_shape[0], dsc_activation, None, False, args.device))
elif args.type == 'ncsn':
model = NetModel(NCSNDde(64, gen_activation, args.device))
else:
print(f'Unknown model type: {args.type}')
if not args.load_from is None:
model.load(args.load_from)
print(f'Loaded model from {args.load_from}')
if len(args.vis_path) > 0:
os.makedirs(args.vis_path, exist_ok=True)
if args.train:
if args.type == 'ncsn':
train_ncsn()
else:
train()
if args.visualize:
visualize(None, args.min_sigma)
BATCH_SIZE = 64
EPOCHS = 20
Z_DIM = 100
D_STEPS = 1
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
t = transforms.Compose([transforms.Resize(64),
transforms.CenterCrop(64),
transforms.ToTensor(),
lambda x: x * 2 - 1]) # Scaling to -1, 1
dataset = torchvision.datasets.CelebA('G:/Datasets', download=DOWNLOAD, transform=t)
dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)
G = Generator(Z_DIM).to(device)
D = Discriminator().to(device)
d_loss = D_loss().to(device)
g_loss = G_loss().to(device)
optim_D = torch.optim.Adam(D.parameters(), lr=1e-5, betas=(0.5, 0.999))
optim_G = torch.optim.Adam(G.parameters(), lr=1e-4, betas=(0.5, 0.999))
d_count = 0
for e in range(EPOCHS):
for x, _ in dataloader:
for p in D.parameters():
p.requires_grad = True
x = x.to(device)
D.zero_grad()
def test_gan(self):
# models settings
lambda_gp = 0.1
Tensor = torch.FloatTensor
generator = Generator(2, 2)
discriminator = Discriminator(2)
optimizer_G = torch.optim.Adam(generator.parameters(),
lr=0.0001,
betas=(0.5, 0.999))
optimizer_D = torch.optim.Adam(discriminator.parameters(),
lr=0.0001,
betas=(0.5, 0.999))
criterion = torch.nn.BCEWithLogitsLoss()
optimizer_D.zero_grad()
# run models for test batch
real_data = Tensor(np.random.normal(0, 1, (64, 2)))
z = Tensor(np.random.normal(0, 1, (64, 2)))
real_target = Tensor(real_data.size(0), 1).fill_(1.0)
fake_target = Tensor(real_data.size(0), 1).fill_(0.0)
g_before = deepcopy(generator)
d_before = deepcopy(discriminator)
# Generate a batch of images
fake_data = generator(z)
# Real images
real_validity = discriminator(real_data)
# Fake images
fake_validity = discriminator(fake_data)
# Gradient penalty
gradient_penalty = compute_gradient_penalty(discriminator,
real_data.data,
fake_data.data, Tensor)
# Discriminator loss
d_loss = criterion(real_validity, real_target) \
+ criterion(fake_validity, fake_target) \
+ lambda_gp * gradient_penalty
d_loss.backward()
optimizer_D.step()
# Assert that D changed and G not changed
g_changed = [
torch.equal(after, before) for after, before in zip(
generator.parameters(), g_before.parameters())
]
self.assertTrue(all(g_changed))
d_changed = [
torch.equal(after, before) for after, before in zip(
discriminator.parameters(), d_before.parameters())
]
self.assertFalse(all(d_changed))
optimizer_G.zero_grad()
# Train on fake samples
g_before = deepcopy(generator)
d_before = deepcopy(discriminator)
fake_data = generator(z)
fake_validity = discriminator(fake_data)
g_loss = criterion(fake_validity, real_target)
g_loss.backward()
optimizer_G.step()
# Assert that G changed and D not changed
g_changed = [
torch.equal(after, before) for after, before in zip(
generator.parameters(), g_before.parameters())
]
self.assertFalse(all(g_changed))
d_changed = [
torch.equal(after, before) for after, before in zip(
discriminator.parameters(), d_before.parameters())
]
self.assertTrue(all(d_changed))
'../data/fashion_mnist/',
train=True,
download=True,
transform=transforms.ToTensor()),
batch_size=64,
shuffle=False,
**kwargs)
test_loader = torch.utils.data.DataLoader(datasets.MNIST(
'../data/fashion_mnist/', train=False, transform=transforms.ToTensor()),
batch_size=32,
shuffle=False,
**kwargs)
test_data = list(test_loader)
netG = Generator().cuda()
netD = Discriminator().cuda()
optimizerG = torch.optim.Adam(netG.parameters(), lr=1e-4, betas=(0.5, 0.9))
optimizerD = torch.optim.Adam(netD.parameters(), lr=1e-4, betas=(0.5, 0.9))
one = torch.cuda.FloatTensor([1])
mone = one * -1
def train(epoch):
train_loss = []
for batch_idx, (data, _) in enumerate(train_loader):
start_time = time.time()
if data.size(0) != 64:
def train(rank: int, cfg: DictConfig):
print(OmegaConf.to_yaml(cfg))
if cfg.train.n_gpu > 1:
init_process_group(backend=cfg.train.dist_config['dist_backend'],
init_method=cfg.train.dist_config['dist_url'],
world_size=cfg.train.dist_config['world_size'] *
cfg.train.n_gpu,
rank=rank)
device = torch.device(
'cuda:{:d}'.format(rank) if torch.cuda.is_available() else 'cpu')
generator = Generator(sum(cfg.model.feature_dims), *cfg.model.cond_dims,
**cfg.model.generator).to(device)
discriminator = Discriminator(**cfg.model.discriminator).to(device)
if rank == 0:
print(generator)
os.makedirs(cfg.train.ckpt_dir, exist_ok=True)
print("checkpoints directory : ", cfg.train.ckpt_dir)
if os.path.isdir(cfg.train.ckpt_dir):
cp_g = scan_checkpoint(cfg.train.ckpt_dir, 'g_')
cp_do = scan_checkpoint(cfg.train.ckpt_dir, 'd_')
steps = 1
if cp_g is None or cp_do is None:
state_dict_do = None
last_epoch = -1
else:
state_dict_g = load_checkpoint(cp_g, device)
state_dict_do = load_checkpoint(cp_do, device)
generator.load_state_dict(state_dict_g['generator'])
discriminator.load_state_dict(state_dict_do['discriminator'])
steps = state_dict_do['steps'] + 1
last_epoch = state_dict_do['epoch']
if cfg.train.n_gpu > 1:
generator = DistributedDataParallel(generator,
device_ids=[rank]).to(device)
discriminator = DistributedDataParallel(discriminator,
device_ids=[rank]).to(device)
optim_g = RAdam(generator.parameters(), cfg.opt.lr, betas=cfg.opt.betas)
optim_d = RAdam(discriminator.parameters(),
cfg.opt.lr,
betas=cfg.opt.betas)
if state_dict_do is not None:
optim_g.load_state_dict(state_dict_do['optim_g'])
optim_d.load_state_dict(state_dict_do['optim_d'])
scheduler_g = torch.optim.lr_scheduler.ExponentialLR(
optim_g, gamma=cfg.opt.lr_decay, last_epoch=last_epoch)
scheduler_d = torch.optim.lr_scheduler.ExponentialLR(
optim_d, gamma=cfg.opt.lr_decay, last_epoch=last_epoch)
train_filelist = load_dataset_filelist(cfg.dataset.train_list)
trainset = FeatureDataset(cfg.dataset, train_filelist, cfg.data)
train_sampler = DistributedSampler(
trainset) if cfg.train.n_gpu > 1 else None
train_loader = DataLoader(trainset,
batch_size=cfg.train.batch_size,
num_workers=cfg.train.num_workers,
shuffle=True,
sampler=train_sampler,
pin_memory=True,
drop_last=True)
if rank == 0:
val_filelist = load_dataset_filelist(cfg.dataset.test_list)
valset = FeatureDataset(cfg.dataset,
val_filelist,
cfg.data,
segmented=False)
val_loader = DataLoader(valset,
batch_size=1,
num_workers=cfg.train.num_workers,
shuffle=False,
sampler=train_sampler,
pin_memory=True)
sw = SummaryWriter(os.path.join(cfg.train.ckpt_dir, 'logs'))
generator.train()
discriminator.train()
for epoch in range(max(0, last_epoch), cfg.train.epochs):
if rank == 0:
start = time.time()
print("Epoch: {}".format(epoch + 1))
if cfg.train.n_gpu > 1:
train_sampler.set_epoch(epoch)
for y, x_noised_features, x_noised_cond in train_loader:
if rank == 0:
start_b = time.time()
y = y.to(device, non_blocking=True)
x_noised_features = x_noised_features.transpose(1, 2).to(
device, non_blocking=True)
x_noised_cond = x_noised_cond.to(device, non_blocking=True)
z1 = torch.randn(cfg.train.batch_size,
cfg.model.cond_dims[1],
device=device)
z2 = torch.randn(cfg.train.batch_size,
cfg.model.cond_dims[1],
device=device)
y_hat1 = generator(x_noised_features, x_noised_cond, z=z1)
y_hat2 = generator(x_noised_features, x_noised_cond, z=z2)
# Discriminator
real_scores, fake_scores = discriminator(y), discriminator(
y_hat1.detach())
d_loss = discriminator_loss(real_scores, fake_scores)
optim_d.zero_grad()
d_loss.backward(retain_graph=True)
optim_d.step()
# Generator
g_stft_loss = criterion(y, y_hat1) + criterion(
y, y_hat2) - criterion(y_hat1, y_hat2)
g_adv_loss = adversarial_loss(fake_scores)
g_loss = g_adv_loss + g_stft_loss
optim_g.zero_grad()
g_loss.backward()
optim_g.step()
if rank == 0:
# STDOUT logging
if steps % cfg.train.stdout_interval == 0:
with torch.no_grad():
print(
'Steps : {:d}, Gen Loss Total : {:4.3f}, STFT Error : {:4.3f}, s/b : {:4.3f}'
.format(steps, g_loss, g_stft_loss,
time.time() - start_b))
# checkpointing
if steps % cfg.train.checkpoint_interval == 0:
ckpt_dir = "{}/g_{:08d}".format(cfg.train.ckpt_dir, steps)
save_checkpoint(
ckpt_dir, {
'generator':
(generator.module if cfg.train.n_gpu > 1 else
generator).state_dict()
})
ckpt_dir = "{}/do_{:08d}".format(cfg.train.ckpt_dir, steps)
save_checkpoint(
ckpt_dir, {
'discriminator':
(discriminator.module if cfg.train.n_gpu > 1 else
discriminator).state_dict(),
'optim_g':
optim_g.state_dict(),
'optim_d':
optim_d.state_dict(),
'steps':
steps,
'epoch':
epoch
})
# Tensorboard summary logging
if steps % cfg.train.summary_interval == 0:
sw.add_scalar("training/gen_loss_total", g_loss, steps)
sw.add_scalar("training/gen_stft_error", g_stft_loss,
steps)
# Validation
if steps % cfg.train.validation_interval == 0:
generator.eval()
torch.cuda.empty_cache()
val_err_tot = 0
with torch.no_grad():
for j, (y, x_noised_features,
x_noised_cond) in enumerate(val_loader):
y_hat = generator(
x_noised_features.transpose(1, 2).to(device),
x_noised_cond.to(device))
val_err_tot += criterion(y, y_hat).item()
if j <= 4:
# sw.add_audio('noised/y_noised_{}'.format(j), y_noised[0], steps, cfg.data.target_sample_rate)
sw.add_audio('generated/y_hat_{}'.format(j),
y_hat[0], steps,
cfg.data.sample_rate)
sw.add_audio('gt/y_{}'.format(j), y[0], steps,
cfg.data.sample_rate)
val_err = val_err_tot / (j + 1)
sw.add_scalar("validation/stft_error", val_err, steps)
generator.train()
steps += 1
scheduler_g.step()
scheduler_d.step()
if rank == 0:
print('Time taken for epoch {} is {} sec\n'.format(
epoch + 1, int(time.time() - start)))
def synth(file_path, domain_A, data_root, output_dir):
sr = 16000
model = Generator(24)
if domain_A:
generator_A2B = torch.load("checkpoint/generator_A2B.pt")
model.load_state_dict(generator_A2B)
else:
generator_B2A = torch.load("checkpoint/generator_B2A.pt")
model.load_state_dict(generator_B2A)
filename_B = os.path.basename(file_path)
SF1_train_data_source = MemoryCacheDataset(
FileSourceDataset((VCC2016DataSource(data_root, ["SF1"],
training=True))))
TF2_train_data_source = MemoryCacheDataset(
FileSourceDataset((VCC2016DataSource(data_root, ["TF2"],
training=True))))
SF1_test_data_source = MemoryCacheDataset(
FileSourceDataset((VCC2016DataSource(data_root, ["SF1"],
training=False))))
TF2_test_data_source = MemoryCacheDataset(
FileSourceDataset((VCC2016DataSource(data_root, ["TF2"],
training=False))))
train_dataset = MCEPWrapper(SF1_train_data_source,
TF2_train_data_source,
mfcc_only=True)
test_dataset = MCEPWrapper(SF1_test_data_source,
TF2_test_data_source,
mfcc_only=False,
norm_calc=False)
test_dataset.input_meanstd = train_dataset.input_meanstd
test_dataset.output_meanstd = train_dataset.output_meanstd
wav, _ = librosa.load(file_path, sr=sr, mono=True)
wav_padded = wav_padding(wav, sr=sr, frame_period=5, multiple=4)
f0, _, sp, ap = world_decompose(wav_padded, sr)
mcep = world_encode_spectral_envelop(sp, sr)
# Normalising MCEPs
mean_A, std_A = train_dataset.input_meanstd
mean_B, std_B = train_dataset.output_meanstd
mean_f0_A = mean_A[0]
mean_f0_B = mean_B[0]
std_f0_A = std_A[0]
std_f0_B = std_B[0]
mean_mcep_A = mean_A[1:25]
mean_mcep_B = mean_B[1:25]
std_mcep_A = std_A[1:25]
std_mcep_B = std_B[1:25]
if domain_A:
normalised_mcep_source = torch.Tensor(
(mcep - mean_mcep_A) / std_mcep_A)
else:
normalised_mcep_source = torch.Tensor(
(mcep - mean_mcep_B) / std_mcep_B)
normalised_mcep_source = normalised_mcep_source[None, :, :]
normalised_mcep_source = normalised_mcep_source.permute(0, 2, 1)
normalised_mcep_target = model(normalised_mcep_source)
normalised_mcep_target = normalised_mcep_target.permute(
0, 2, 1).cpu().detach().numpy()[0, :, :]
if domain_A:
mcep_target = normalised_mcep_target * std_mcep_B + mean_mcep_B
else:
mcep_target = normalised_mcep_target * std_mcep_B + mean_mcep_B
mcep_target = np.ascontiguousarray(mcep_target)
# Because here we directly decompose from signal, the nonmasked array implementation is used
f0_target = pitch_conversion(f0, mean_f0_A, std_f0_A, mean_f0_B, std_f0_B)
sp_target = world_decode_spectral_envelop(mcep_target, sr)
ap_target = np.ascontiguousarray(ap)
speech_fake_A = world_speech_synthesis(f0_target,
sp_target,
ap_target,
sr,
frame_period=5)
librosa.output.write_wav(os.path.join(output_dir, filename_B),
speech_fake_A, sr)
v for k, v in os.environ.items() if k.startswith('GEN_SRC')
]
for v in check_paths:
if m := re_path.match(v):
if not os.path.exists(m.group('path')):
print(f'Source directory cannot be found: {v}')
exit(0)
src_paths[m.group('type')] = m.group('path')
kwargs['src_paths'] = src_paths
if not kwargs.get('pg_uri'):
raise RuntimeError('PostgreSQL URI is required')
elif not kwargs.get('src_paths'):
raise RuntimeError('Source path(s) required')
elif not kwargs.get('out_path'):
raise RuntimeError('Output path required')
inst = Generator(LOOP, **kwargs)
async def _run():
await inst.init_task
await inst.replace_all()
try:
await inst.check_new_files()
except KeyboardInterrupt:
inst.logger.info("Keyboard interrupt, exit.")
except Exception as error:
traceback.print_exception(type(error), error, error.__traceback__)
await inst.close()
LOOP.run_until_complete(_run())
parser.add_argument("--n_critic", type=int, default=1, help="number of training steps for discriminator per iter")
opt = parser.parse_args()
# Write learning info
writer = SummaryWriter('runs/2d_gaussain_mixture')
np.random.seed(opt.random_seed)
torch.manual_seed(opt.random_seed)
cuda = True if torch.cuda.is_available() and opt.cuda else False
Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() and opt.cuda else torch.FloatTensor
# Loss weight for gradient penalty
lambda_gp = opt.lambda_gp
# Initialize generator and discriminator
generator = Generator(opt.latent_dim, 2)
discriminator = Discriminator(2)
gaussian = GaussianMixture(opt.grid_side_size)
train_data = gaussian.sample(opt.batch_size * opt.target_factor if opt.learning_type == 'smallgan' else opt.batch_size)
if cuda:
generator.cuda()
discriminator.cuda()
prior_batch_size = opt.batch_size * opt.prior_factor if opt.learning_type == 'smallgan' else opt.batch_size
# Valid dataset
valid_samples = np.random.uniform(-1, 1, (opt.valid_size, opt.latent_dim))
valid_dataloader = torch.utils.data.DataLoader(
valid_samples,
batch_size=opt.valid_batch_size,
def generator(cfg, n_features, cond_dim, z_dim):
generator = Generator(n_features, cond_dim, z_dim, **cfg.model.generator)
return generator
mfcc_only=True)
test_dataset = MCEPWrapper(SF1_test_data_source,
TF2_test_data_source,
mfcc_only=False,
norm_calc=False)
test_dataset.input_meanstd = train_dataset.input_meanstd
test_dataset.output_meanstd = train_dataset.output_meanstd
train_dataset_loader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True)
test_dataset_loader = DataLoader(test_dataset,
batch_size=batch_size,
shuffle=False)
generator_A2B = Generator(24).to("cuda")
generator_B2A = Generator(24).to("cuda")
discriminator_A = Discriminator(1).to("cuda")
discriminator_B = Discriminator(1).to("cuda")
generator_params = [generator_A2B.parameters(), generator_B2A.parameters()]
generator_optimizer = torch.optim.Adam(itertools.chain(*generator_params),
lr=generator_lr)
discriminator_params = [
discriminator_A.parameters(),
discriminator_B.parameters()
]
discriminator_optimizer = torch.optim.Adam(
itertools.chain(*discriminator_params), lr=discriminator_lr)
for epoch in range(num_epochs):
import torch
from torch.utils.data import DataLoader
import torchvision
import torchvision.transforms as transforms
import numpy as np
from modules import Generator, Discriminator
from wgangp import D_loss, G_loss, calc_gradient_penalty
import sys
device = torch.device('cuda') if torch.cuda.is_available() else torch.device(
'cpu')
G = Generator().to(device)
G.load_state_dict(torch.load('G.nn'))
with torch.no_grad():
z = torch.randn((64, 100)).to(device)
Gz = G(z)
Gz = (Gz + 1) / 2
torchvision.utils.save_image(Gz, f'./sample.jpg')
