idx,
dataformats='HW')
writer.flush()
exit()
# -----------------------------------
embedding = nn.Embedding(num_embeddings=len(_symbol_to_id), embedding_dim=512)
encoder = Encoder(emb_channels=512, enc_out=512).to(device)
decoder = Decoder(mel_channels=80).to(device)
optimizer = torch.optim.Adam([{
'params': embedding.parameters()
}, {
'params': encoder.parameters()
}, {
'params': decoder.parameters()
}],
lr=0.001,
betas=(0., 0.999))
# -----------------------------------
global_idx = 0
mean_losses = np.zeros(3)
for epoch in enumerate(tqdm(dataloader)):
for idx, batch in enumerate(tqdm(dataloader)):
text_data, text_len, text_mask, mel_data, mel_len, mel_mask = batch
# audio_data = F.avg_pool1d(audio_data, kernel_size=2, padding=1)
text_emb = embedding(text_data)
def main():
global args
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
dataset_base_path = path.join(args.base_path, "dataset", "celeba")
image_base_path = path.join(dataset_base_path, "img_align_celeba")
split_dataset_path = path.join(dataset_base_path, "Eval", "list_eval_partition.txt")
with open(split_dataset_path, "r") as f:
split_annotation = f.read().splitlines()
# create the data name list for train,test and valid
train_data_name_list = []
test_data_name_list = []
valid_data_name_list = []
for item in split_annotation:
item = item.split(" ")
if item[1] == '0':
train_data_name_list.append(item[0])
elif item[1] == '1':
valid_data_name_list.append(item[0])
else:
test_data_name_list.append(item[0])
attribute_annotation_dict = None
if args.need_label:
attribute_annotation_path = path.join(dataset_base_path, "Anno", "list_attr_celeba.txt")
with open(attribute_annotation_path, "r") as f:
attribute_annotation = f.read().splitlines()
attribute_annotation = attribute_annotation[2:]
attribute_annotation_dict = {}
for item in attribute_annotation:
img_name, attribute = item.split(" ", 1)
attribute = tuple([eval(attr) for attr in attribute.split(" ") if attr != ""])
assert len(attribute) == 40, "the attribute of item {} is not equal to 40".format(img_name)
attribute_annotation_dict[img_name] = attribute
discriminator = Discriminator(num_channel=args.num_channel, num_feature=args.dnf,
data_parallel=args.data_parallel).cuda()
generator = Decoder(latent_dim=args.latent_dim, num_feature=args.gnf,
num_channel=args.num_channel, data_parallel=args.data_parallel).cuda()
input("Begin the {} time's training, the train dataset has {} images and the valid has {} images".format(
args.train_time, len(train_data_name_list), len(valid_data_name_list)))
d_optimizer = torch.optim.Adam(discriminator.parameters(), lr=args.lr, betas=(args.beta1, 0.999))
g_optimizer = torch.optim.Adam(generator.parameters(), lr=args.lr, betas=(args.beta1, 0.999))
writer_log_dir = "{}/DCGAN/runs/train_time:{}".format(args.base_path, args.train_time)
# Here we implement the resume part
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
if not args.not_resume_arg:
args = checkpoint['args']
args.start_epoch = checkpoint['epoch']
discriminator.load_state_dict(checkpoint["discriminator_state_dict"])
generator.load_state_dict(checkpoint["generator_state_dict"])
d_optimizer.load_state_dict(checkpoint['discriminator_optimizer'])
g_optimizer.load_state_dict(checkpoint['generator_optimizer'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
else:
raise FileNotFoundError("Checkpoint Resume File {} Not Found".format(args.resume))
else:
if os.path.exists(writer_log_dir):
flag = input("DCGAN train_time:{} will be removed, input yes to continue:".format(
args.train_time))
if flag == "yes":
shutil.rmtree(writer_log_dir, ignore_errors=True)
writer = SummaryWriter(log_dir=writer_log_dir)
# Here we just use the train dset in training
train_dset = CelebADataset(base_path=image_base_path, data_name_list=train_data_name_list,
image_size=args.image_size,
label_dict=attribute_annotation_dict)
train_dloader = DataLoader(dataset=train_dset, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True)
criterion = nn.BCELoss()
for epoch in range(args.start_epoch, args.epochs):
train(train_dloader, generator, discriminator, g_optimizer, d_optimizer, criterion, writer, epoch)
# save parameters
save_checkpoint({
'epoch': epoch + 1,
'args': args,
"discriminator_state_dict": discriminator.state_dict(),
"generator_state_dict": generator.state_dict(),
'discriminator_optimizer': d_optimizer.state_dict(),
'generator_optimizer': g_optimizer.state_dict()
})
