def get_dataset_props(data_dir):
big_data_dir = os.path.join(cfg.DATA_DIR, 'big')
captions_path = os.path.join(cfg.DATA_DIR, cfg.DATA_SIZE,
'captions.pickle')
delete_file(captions_path)
imsize = cfg.TREE.BASE_SIZE * (2**(cfg.TREE.BRANCH_NUM - 1))
image_transform = transforms.Compose([
transforms.Resize(int(imsize * 76 / 64)),
transforms.RandomCrop(imsize),
transforms.RandomHorizontalFlip()
])
dataset = TextDataset(big_data_dir,
'test',
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
n_words = dataset.n_words
wordtoix = dataset.wordtoix
delete_file(captions_path)
dataset = TextDataset(data_dir,
'test',
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
captions = []
for filename in dataset.filenames:
filepath = '%s/text/%s.txt' % (data_dir, filename)
with open(filepath, "r", encoding="utf-8") as f:
caps = f.read().split('\n')
cap = ''
trials = 0
# Take a random caption
while cap == '' or trials > cfg.TEXT.CAPTIONS_PER_IMAGE:
if caps[trials] != '':
cap = caps[trials]
trials += 1
captions.append(cap)
delete_file(captions_path)
counter = 0
for cap in captions:
if cap == '':
counter += 1
return n_words, wordtoix, captions
def draw_img():
args = parse_args()
if args.cfg_file is not None:
cfg_from_file(args.cfg_file)
if args.data_dir != '':
cfg.DATA_DIR = args.data_dir
args.manualSeed = random.randint(1, 10000)
random.seed(args.manualSeed)
np.random.seed(args.manualSeed)
torch.manual_seed(args.manualSeed)
if cfg.CUDA:
torch.cuda.manual_seed_all(args.manualSeed)
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = True
print("Seed: %d" % (args.manualSeed))
special_model_path = 'Model/coco2/coco2.pth'
dataset = TextDataset(cfg.DATA_DIR, base_size=cfg.TREE.BASE_SIZE)
assert dataset
gen_example(dataset.n_words, dataset.wordtoix, dataset.ixtoword, special_model_path)
def test():
cfgs = Cfgs()
imsize = cfgs.image_size
DATA_DIR = "../data/coco"
image_transform = transforms.Compose([
transforms.Resize(int(imsize * 76 / 64)), #]) #,
transforms.RandomCrop(imsize),
transforms.RandomHorizontalFlip()
])
ts = transforms.Compose([transforms.Resize(imsize), transforms.ToTensor()])
dataset = TextDataset(DATA_DIR,
'train',
base_size=64,
transform=image_transform,
target_transform=ts)
print("dataset.n_words:", dataset.n_words, "dataset.embeddings_num:",
dataset.embeddings_num)
dataloader = DataLoader(dataset,
batch_size=cfgs.batch_size,
drop_last=True,
shuffle=True,
num_workers=5)
print(DATA_DIR)
print("****")
for batch in dataloader:
print(type(batch))
def __init__(self, eval_ = False,
inn_channels = cfg.channels,
generatorLR = cfg.generatorLR,
discriminatorLR = cfg.discriminatorLR,
StageNum = 4, beta1 = 0.5, beta2 = 0.999,
zDim = 100):
self.zDim = zDim
self.inn_channels = inn_channels
self.eval_ = eval_
self.StageNum = StageNum
if cfg.inception:
from losses import INCEPTION_V3, compute_inception_score
self.inception_model = INCEPTION_V3()
inception_model = inception_model.cuda()
inception_model.eval()
image_transform = transforms.Compose([
transforms.Scale(int((64 * 4) * 76 / 64)),
transforms.RandomCrop(64 * 4),
transforms.RandomHorizontalFlip()])
self.dataset = TextDataset('birds', 'train',
base_size=64,
transform=image_transform,
StageNum=StageNum)
self.generator = models.G_NET(StageNum=StageNum, zDim=zDim).cuda()
self.discriminator = []
if StageNum == 1:
self.discriminator.append(D_NET64(self.inn_channels).cuda())
elif StageNum == 2:
self.discriminator.append(D64(self.inn_channels).cuda())
self.discriminator.append(D128(self.inn_channels).cuda())
elif StageNum == 3:
self.discriminator.append(D64(self.inn_channels).cuda())
self.discriminator.append(D128(self.inn_channels).cuda())
self.discriminator.append(D256(self.inn_channels).cuda())
elif StageNum == 4:
self.discriminator.append(D64(self.inn_channels).cuda())
self.discriminator.append(D128(self.inn_channels).cuda())
self.discriminator.append(D256(self.inn_channels).cuda())
self.discriminator.append(D512(self.inn_channels).cuda())
elif StageNum == 5:
self.discriminator.append(D64(self.inn_channels).cuda())
self.discriminator.append(D128(self.inn_channels).cuda())
self.discriminator.append(D256(self.inn_channels).cuda())
self.discriminator.append(D512(self.inn_channels).cuda())
self.discriminator.append(D1024(self.inn_channels).cuda())
self.generator.apply(models.weights_init)
for i in range(len(self.discriminator)):
self.discriminator[i].apply(models.weights_init)
self.loss = torch.nn.BCELoss().cuda()
from torch.optim import Adam
self.gOptimizer = Adam(self.generator.parameters(), lr=generatorLR, betas=(beta1, beta2))
self.disOptimizer = []
for i in range(len(self.discriminator)):
opt = Adam(self.discriminator[i].parameters(), lr=discriminatorLR, betas=(beta1, beta2))
self.disOptimizer.append(opt)
def main():
start_t = time.time()
seed = 12345
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if CONFIG['CUDA']:
torch.cuda.manual_seed_all(seed)
now = datetime.datetime.now(dateutil.tz.tzlocal())
timestamp = now.strftime('%Y_%m_%d_%H_%M_%S')
output_dir = '../output/%s_%s_%s' % \
(CONFIG['DATASET_NAME'], CONFIG['CONFIG_NAME'], timestamp)
split_dir, bshuffle = 'train', True
if not CONFIG['TRAIN']['FLAG']:
# bshuffle = False
split_dir = 'test'
imsize = CONFIG['TREE']['BASE_SIZE'] * (2**(CONFIG['TREE']['BRANCH_NUM'] -
1))
image_transform = transforms.Compose([
transforms.Scale(int(imsize * 76 / 64)),
transforms.RandomCrop(imsize),
transforms.RandomHorizontalFlip()
])
dataset = TextDataset(CONFIG['DATA_DIR'],
split_dir,
base_size=CONFIG['TREE']['BASE_SIZE'],
transform=image_transform)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=CONFIG['TRAIN']['BATCH_SIZE'],
drop_last=True,
shuffle=bshuffle,
num_workers=int(CONFIG['WORKERS']))
algo = trainer(output_dir, dataloader, dataset.n_words, dataset.ixtoword)
if CONFIG['TRAIN']['FLAG']:
algo.train()
else:
'''generate images from pre-extracted embeddings'''
if CONFIG['B_VALIDATION']:
algo.sampling(
split_dir) # generate images for the whole valid dataset
else:
gen_example(dataset.wordtoix,
algo) # generate images for customized /
end_t = time.time()
print('Total time for training:', end_t - start_t)
def job(confidence, start, end, order):
# use multiple process to accelerate procedure
print("Thread order: %d" % order)
print("Work lower: %d, upper: %d" % (start, end))
class_descriptions = pd.read_csv('../data/OpenImage/class-descriptions.csv', header=None)
ixtolabel = {}
for i in range(class_descriptions.shape[0]):
ixtolabel[class_descriptions.iloc[i][0]] = class_descriptions.iloc[i][1].lower()
print("Find Same Labels preprocessing! Strart count time!")
t0 = time.clock()
dataset = TextDataset('../data/coco', 'train', base_size=299)
wordtoix = dataset.wordtoix
df_labels = pd.read_csv('../data/OpenImage/class-descriptions.csv', header=None)
labels = {}
id_labels = {}
for i in range(df_labels.shape[0]):
labels[df_labels.iloc[i][1].lower()] = df_labels.iloc[i][0]
id_labels[df_labels.iloc[i][0]] = df_labels.iloc[i][1].lower()
#print(labels.keys())
same_labels = []
counts = 0
for key in list(labels.keys()):
if key in wordtoix:
same_labels.append(key)
counts += 1
print(counts)
machine_label_path = 'train-annotations-machine-imagelabels.csv'
# del df_labels
df_machine_labels = pd.read_csv('../data/OpenImage/'+machine_label_path)
image_dict = {}
# print('Thread order: %d, Finish preprocess same labels, Used time: %.2f' % (order, time.clock()-t0))
print('\nThread order: %d, Start go through machine-label data! Start count time!' % order)
t0 = time.time()
for i in range(start, end):
if ((i-start) % 400000) == 0 and (i-start) != 0:
print('Thread order: %d, Now %.2f%% Completed !, Used time: %.2f' % (order, ((i-start)/(end-start))*100., time.time()-t0))
t0 = time.time()
label = ixtolabel[df_machine_labels.iloc[i]['LabelName']]
if df_machine_labels.iloc[i]['ImageID'] not in image_dict and df_machine_labels.iloc[i]['Confidence'] >= confidence:
image_dict[df_machine_labels.iloc[i]['ImageID']] = []
image_dict[df_machine_labels.iloc[i]['ImageID']].append(label)
elif df_machine_labels.iloc[i]['Confidence'] >= confidence:
image_dict[df_machine_labels.iloc[i]['ImageID']].append(label)
print('Thread order: %d, Finsih machine-label, Total Image: %d' % (order, len(list(image_dict.keys()))))
save_path = '../data/machine-data' + str(order) + '.pickle'
with open(save_path, 'wb') as f:
pickle.dump(image_dict, f, pickle.HIGHEST_PROTOCOL)
def draw_picture(label):
if label == 'bird':
cfg_from_file('sample.yml')
special_model_path = 'Model/bird_2.pth'
special_model_path = 'Model/netG_epoch_575.pth'
elif label == 'coco':
cfg_from_file('sample_coco.yml')
special_model_path = 'Model/coco/coco_2.pth'
#special_model_path = 'Model/coco/netG_epoch_55.pth'
manualSeed = random.randint(1, 10000)
random.seed(manualSeed)
np.random.seed(manualSeed)
torch.manual_seed(manualSeed)
print('generating images for customized captions ... ')
dataset = TextDataset(cfg.DATA_DIR, base_size=cfg.TREE.BASE_SIZE)
assert dataset
# print(dataset.n_words)
if label == 'bird' or 'coco':
gen_example(dataset.n_words, dataset.wordtoix, dataset.ixtoword,
special_model_path)
eval = True
# Get data loader
imsize = cfg.TREE.BASE_SIZE * (2**(cfg.TREE.BRANCH_NUM - 1))
image_transform = transforms.Compose(
[transforms.Resize((268, 268)),
transforms.ToTensor()])
if cfg.TRAIN.OPTIMIZE_DATA_LOADING:
num_max_objects = 10
dataset_indices = get_dataset_indices(num_max_objects=num_max_objects, split="train"\
if cfg.TRAIN.FLAG else "test")
dataset = TextDataset(cfg.DATA_DIR,
img_dir,
split_dir,
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform,
eval=eval,
use_generated_bboxes=cfg.TRAIN.GENERATED_BBOXES)
assert dataset
dataset_subsets = []
dataloaders = []
for max_objects in range(num_max_objects + 1):
subset = torch.utils.data.Subset(dataset,
dataset_indices[max_objects])
print(subset)
dataset_subsets.append(subset)
print(int(cfg.WORKERS))
dataloader = torch.utils.data.DataLoader(
subset,
batch_size=cfg.TRAIN.BATCH_SIZE[max_objects],
split_dir, bshuffle = 'train', True
if not cfg.TRAIN.FLAG:
# bshuffle = False
split_dir = 'test'
# Get data loader
imsize = cfg.TREE.BASE_SIZE * (2**(cfg.TREE.BRANCH_NUM - 1))
image_transform = transforms.Compose([
transforms.Scale(int(imsize * 76 / 64)),
transforms.RandomCrop(imsize),
transforms.RandomHorizontalFlip()
])
dataset = TextDataset(cfg.DATA_DIR,
split_dir,
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform,
max_length=config.max_position_embeddings,
vocab=config.vocab)
assert dataset
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=cfg.TRAIN.BATCH_SIZE,
drop_last=True,
shuffle=bshuffle,
num_workers=int(cfg.WORKERS))
dataset_val = TextDataset(cfg.DATA_DIR,
'test',
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
val_batch_size = 24
dataloader_val = torch.utils.data.DataLoader(dataset_val,
def main(gpu_id, data_dir, manual_seed, cuda, train_flag, image_size,
batch_size, workers, stage, dataset_name, config_name, max_epoch,
snapshot_interval, net_g, net_d, z_dim, generator_lr,
discriminator_lr, lr_decay_epoch, coef_kl, stage1_g, embedding_type,
condition_dim, df_dim, gf_dim, res_num, text_dim, regularizer):
if manual_seed is None:
manual_seed = random.randint(1, 10000)
random.seed(manual_seed)
torch.manual_seed(manual_seed)
if cuda:
torch.cuda.manual_seed_all(manual_seed)
now = datetime.datetime.now(dateutil.tz.tzlocal())
timestamp = now.strftime('%Y_%m_%d_%H_%M_%S')
output_dir = '../output/%s_%s_%s' % (dataset_name, config_name, timestamp)
num_gpu = len(gpu_id.split(','))
if train_flag:
image_transform = transforms.Compose([
transforms.RandomCrop(image_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataset = TextDataset(data_dir,
'train',
imsize=image_size,
transform=image_transform)
assert dataset
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size *
num_gpu,
drop_last=True,
shuffle=True,
num_workers=int(workers))
algo = GANTrainer(output_dir, max_epoch, snapshot_interval, gpu_id,
batch_size, train_flag, net_g, net_d, cuda, stage1_g,
z_dim, generator_lr, discriminator_lr,
lr_decay_epoch, coef_kl, regularizer)
algo.train(dataloader, stage, text_dim, gf_dim, condition_dim, z_dim,
df_dim, res_num)
elif dataset_name == 'birds' and train_flag is False:
image_transform = transforms.Compose([
transforms.RandomCrop(image_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataset = TextDataset(data_dir,
'train',
imsize=image_size,
transform=image_transform)
assert dataset
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size *
num_gpu,
drop_last=True,
shuffle=True,
num_workers=int(workers))
algo = GANTrainer(output_dir, max_epoch, snapshot_interval, gpu_id,
batch_size, train_flag, net_g, net_d, cuda, stage1_g,
z_dim, generator_lr, discriminator_lr,
lr_decay_epoch, coef_kl, regularizer)
algo.birds_eval(dataloader, stage)
else:
datapath = '%s/test/val_captions.t7' % (data_dir)
algo = GANTrainer(output_dir, max_epoch, snapshot_interval, gpu_id,
batch_size, train_flag, net_g, net_d, cuda, stage1_g,
z_dim, generator_lr, discriminator_lr,
lr_decay_epoch, coef_kl, regularizer)
algo.sample(datapath, stage)
##########################################################################
torch.cuda.set_device(cfg.GPU_ID)
cudnn.benchmark = True
# Get data loader ##################################################
imsize = cfg.TREE.BASE_SIZE * (2**(cfg.TREE.BRANCH_NUM - 1))
batch_size = cfg.TRAIN.BATCH_SIZE
image_transform = transforms.Compose([
transforms.Resize(int(imsize * 76 / 64)),
transforms.RandomCrop(imsize),
transforms.RandomHorizontalFlip()
])
valid_dir = args.Vdir
dataset_val = TextDataset(cfg.DATA_DIR,
'test',
valid_dir,
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
dataloader_val = torch.utils.data.DataLoader(dataset_val,
batch_size=batch_size,
drop_last=True,
shuffle=True,
num_workers=int(cfg.WORKERS))
# calculate R precision ##############################################################
Inception_encoder, labels = build_models()
netT = torch.load(cfg.TRAIN.NET_T)
text_encoder = torch.load(cfg.TRAIN.NET_E)
netT = netT.cuda()
text_encoder = text_encoder.cuda()
mkdir_p(folder)
fullpath = os.path.join(save_dir, name) + ".png"
# range from [-1, 1] to [0, 255]
img = images[i].add(1).div(2).mul(255).clamp(0, 255).byte()
ndarr = img.permute(1, 2, 0).data.cpu().numpy()
im = Image.fromarray(ndarr)
im.save(fullpath)
DATA_DIR= 'C:\\Users\\Alper\\PycharmProjects\\MSGAN\\datasets\\birds'
imsize = 256
image_transform = transforms.Compose([
transforms.Resize(int(imsize * 76 / 64)),
transforms.RandomCrop(imsize),
transforms.RandomHorizontalFlip()])
dataset = TextDataset(DATA_DIR, split='test', base_size=64,
transform=image_transform)
nz = 100
n_samples = 10
fid_model = FID_INCEPTION()
fid_model.cuda()
fid_model.eval()
inception_model = INCEPTION_V3()
inception_model.cuda()
inception_model.eval()
G_NET_Path = 'C:\\Users\\alper\\PycharmProjects\\MSGAN\\StackGAN++-Mode-Seeking\\models\\ours_new.pth'
netG = G_NET()
netG.apply(weights_init)
torch.cuda.set_device(cfg.GPU_ID)
cudnn.benchmark = True
# Get data loader ##################################################
imsize = cfg.TREE.BASE_SIZE * (2**(cfg.TREE.BRANCH_NUM - 1))
batch_size = cfg.TRAIN.BATCH_SIZE
# image_transform = transforms.Compose([
# transforms.Scale(int(imsize * 76 / 64)),
# transforms.RandomCrop(imsize),
# transforms.RandomHorizontalFlip()])
image_transform = transforms.Compose(
[transforms.Resize(imsize),
transforms.CenterCrop(imsize)])
dataset = TextDataset(cfg.DATA_DIR,
'train',
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform,
input_channels=args.input_channels,
image_type=args.image_type)
# print(dataset.ixtoword)
# assert False
print(dataset.n_words, dataset.embeddings_num)
assert dataset
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
drop_last=True,
shuffle=False,
num_workers=int(cfg.WORKERS))
if __name__ == "__main__":
args = parse_args()
if args.cfg_file is not None:
cfg_from_file(args.cfg_file)
if args.data_dir != '':
cfg.DATA_DIR = args.data_dir
print('Using config:')
pprint.pprint(cfg)
args.manualSeed = random.randint(1, 10000)
random.seed(args.manualSeed)
np.random.seed(args.manualSeed)
torch.manual_seed(args.manualSeed)
if cfg.CUDA:
torch.cuda.manual_seed_all(args.manualSeed)
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = True
print("Seed: %d" % (args.manualSeed))
special_model_path = 'Model/coco2/coco2.pth'
dataset = TextDataset(cfg.DATA_DIR, base_size=cfg.TREE.BASE_SIZE)
assert dataset
gen_example(dataset.n_words, dataset.wordtoix, dataset.ixtoword, special_model_path)
print("done!")
def pretrain_STREAM():
# rasnet tranformation/normalization
transform = transforms.Compose(
[transforms.RandomCrop(224),
transforms.RandomHorizontalFlip()])
norm = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
now = datetime.now(dateutil.tz.tzlocal())
timestamp = now.strftime('%Y_%m_%d_%H_%M_%S')
output_dir = os.path.join(cfg.OUTPUT_PATH, cfg.DATASET_NAME,
cfg.CONFIG_NAME, timestamp)
mkdir_p(output_dir)
######################
# Run training/validation
######################
crit = nn.CrossEntropyLoss().to(cfg.DEVICE)
# Get data loader ##################################################
imsize = cfg.TREE.BASE_SIZE * (2**(cfg.TREE.BRANCH_NUM - 1))
batch_size = cfg.TRAIN.BATCH_SIZE
image_transform = transforms.Compose([
transforms.Resize(int(imsize * 76 / 64)),
transforms.RandomCrop(imsize),
transforms.RandomHorizontalFlip()
])
split = 'train' if cfg.TRAIN.FLAG else 'test'
dataset = TextDataset(cfg.DATA_DIR,
split,
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
print(dataset.n_words, dataset.embeddings_num)
assert dataset
caption_cnn, caption_rnn, dec_optim = init_model(dataset.ixtoword)
# Load data
dataloader = torch.utils.data.DataLoader(dataset,
batch_size,
drop_last=True,
shuffle=True,
num_workers=int(cfg.WORKERS))
if cfg.TRAIN.FLAG:
train(caption_cnn,
caption_rnn,
decoder_optimizer=dec_optim,
criterion=crit,
train_loader=dataloader,
output_dir=output_dir)
else:
caption_cnn.eval()
caption_rnn.eval()
# Don't caluclate gradients for validation
with torch.no_grad():
validate(caption_cnn,
caption_rnn,
val_loader=dataloader,
ixtoword=dataset.ixtoword)
else:
new_word.append(new_char)
# t_char = dataset.to_one_hot(new_char)
t_word = torch.cat([t_word, char_idx.view(1, 1)], dim=1)
return new_word
if __name__ == '__main__':
data_root = '/Users/charlesren/Downloads/AI_lab_data'
max_length = 6
# Prepare GPU.
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# Prepare dataset.
dataset = TextDataset(data_root, max_length)
input_size = 68
hidden_size = 16
# Prepare model.
model = LSTMAE(input_size, hidden_size, 1, False)
model.load_state_dict(torch.load('model.pt'))
model = model.to(device)
new_words = []
for letter in 'Metallen':
word = generate(letter, dataset, model)
print(word)
new_words.append(word)
mkdir_p(image_dir)
torch.cuda.set_device(cfg.GPU_ID)
cudnn.benchmark = True
# Get data loader ##################################################
imsize = cfg.TREE.BASE_SIZE * (2**(cfg.TREE.BRANCH_NUM - 1))
batch_size = cfg.TRAIN.BATCH_SIZE
image_transform = transforms.Compose([
transforms.Resize(int(imsize * 76 / 64)),
transforms.RandomCrop(imsize),
transforms.RandomHorizontalFlip()
])
dataset = TextDataset(cfg.DATA_DIR,
args.text_encoder_type,
'train',
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
print(dataset.n_words, dataset.embeddings_num)
assert dataset
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
drop_last=True,
shuffle=True,
num_workers=int(cfg.WORKERS))
# # validation data #
print(args.text_encoder_type)
dataset_val = TextDataset(cfg.DATA_DIR,
args.text_encoder_type,
split_dir, bshuffle = 'train', True
if not cfg.TRAIN.FLAG:
# bshuffle = False
split_dir = 'test'
# Get data loader
# TODO: Make all data setup a separate module
imsize = cfg.TREE.BASE_SIZE * (2**(cfg.TREE.BRANCH_NUM - 1))
image_transform = transforms.Compose([
transforms.Resize(int(imsize * 76 / 64)),
transforms.RandomCrop(imsize),
transforms.RandomHorizontalFlip()
])
dataset = TextDataset(cfg.DATA_DIR,
args.text_encoder_type,
split_dir,
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
assert dataset
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=cfg.TRAIN.BATCH_SIZE,
drop_last=True,
shuffle=bshuffle,
num_workers=int(cfg.WORKERS))
start_t = time.time()
if cfg.TRAIN.FLAG:
# Define trainer
trainer = condGANTrainer(output_dir, dataloader, dataset.n_words,
dataset.ixtoword, dataset.text_encoder_type)
# Train
(cfg.DATA_DIR, cfg.DATASET_NAME, split_dir),
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
elif cfg.DATA_DIR.find('imagenet') != -1:
from datasets import ImageFolder
dataset = ImageFolder(cfg.DATA_DIR,
split_dir='train',
custom_classes=CLASS_DIC[cfg.DATASET_NAME],
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
elif cfg.GAN.B_CONDITION: # text to image task
from datasets import TextDataset
if cfg.DATASET_NAME == 'zappos':
dataset = TextDataset(cfg.DATA_DIR,
split_dir,
base_size=cfg.TREE.BASE_SIZE,
embedding_type=cfg.EMBEDDING_TYPE,
transform=image_transform)
else:
dataset = TextDataset(cfg.DATA_DIR,
split_dir,
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
assert dataset
num_gpu = len(cfg.GPU_ID.split(','))
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=cfg.TRAIN.BATCH_SIZE *
num_gpu,
drop_last=True,
shuffle=bshuffle,
num_workers=int(cfg.WORKERS))
folds = GroupKFold(n_splits=n_folds).split(X=train['question_body'],
groups=train['question_body'])
oofs = np.zeros((len(train), N_TARGETS))
main_logger.info(f'Start training model {model_name}...')
for fold_id, (train_index, valid_index) in enumerate(folds):
main_logger.info(f'Fold {fold_id + 1} started at {time.ctime()}')
fold_logger = init_logger(log_dir,
f'train_fold_{fold_id+1}_{model_name}.log')
train_loader = DataLoader(
TextDataset(cat_features_train, ids_train['question'],
ids_train['answer'], seg_ids_train['question'],
seg_ids_train['answer'], train_index, y),
batch_size=bs,
shuffle=True,
num_workers=num_workers)
valid_loader = DataLoader(
TextDataset(cat_features_train, ids_train['question'],
ids_train['answer'], seg_ids_train['question'],
seg_ids_train['answer'], valid_index, y),
batch_size=bs,
shuffle=False,
num_workers=num_workers)
model = models[model_name]()
optimizer = get_optimizer(model, lr, weight_decay, model_type)
folds = GroupKFold(n_splits=n_folds).split(X=train['question_body'],
groups=train['question_body'])
oofs = np.zeros((len(train), N_TARGETS))
main_logger.info(f'Start finetuning model {model_name}...')
for fold_id, (train_index, valid_index) in enumerate(folds):
main_logger.info(f'Fold {fold_id + 1} started at {time.ctime()}')
fold_logger = init_logger(
log_dir, f'finetune_fold_{fold_id+1}_{model_name}.log')
loader = DataLoader(TextDataset(cat_features_train,
ids_train['question'],
ids_train['answer'],
seg_ids_train['question'],
seg_ids_train['answer'],
np.arange(len(train)), y),
batch_size=bs,
shuffle=False,
num_workers=num_workers)
model = models[model_name]()
checkpoint_file = f'{checkpoint_dir}{model_name}_fold_{fold_id+1}_best.pth'
# Get last hidden layer outputs from transformers
fold_logger.info(
f'Precompute transformer outputs for model {model_name}...')
q_outputs, a_outputs = get_model_outputs(model, loader,
checkpoint_file, device,
model_type)
def loading_model(dataset_name='bird'):
#IMPORTANT ARGUMENTS
if (dataset_name=='bird') :
cfg_file=os.path.join(current_dir,"cfg/eval_bird.yml")
else :
cfg_file=os.path.join(current_dir,"cfg/eval_coco.yml")
gpu_id=-1 #change it to 0 or more when using gpu
data_dir=''
manualSeed = 100
#cfg file set
if cfg_file is not None:
cfg_from_file(cfg_file)
if gpu_id != -1:
cfg.GPU_ID = gpu_id
else:
cfg.CUDA = False
if data_dir != '':
cfg.DATA_DIR = data_dir
now = datetime.datetime.now(dateutil.tz.tzlocal())
timestamp = now.strftime('%Y_%m_%d_%H_%M_%S')
output_dir = '../output/%s_%s_%s' % \
(cfg.DATASET_NAME, cfg.CONFIG_NAME, timestamp)
split_dir, bshuffle = 'train', True
if not cfg.TRAIN.FLAG:
# bshuffle = False
split_dir = 'test'
# Get data loader
imsize = cfg.TREE.BASE_SIZE * (2 ** (cfg.TREE.BRANCH_NUM - 1))
image_transform = transforms.Compose([
transforms.Scale(int(imsize * 76 / 64)),
transforms.RandomCrop(imsize),
transforms.RandomHorizontalFlip()])
dataset = TextDataset(cfg.DATA_DIR, split_dir,
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
assert dataset
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=cfg.TRAIN.BATCH_SIZE,
drop_last=True, shuffle=bshuffle, num_workers=int(cfg.WORKERS))
###setting up ALGO
# Define models and go to train/evaluate
algo = trainer(output_dir, dataloader, dataset.n_words, dataset.ixtoword)
#loading text ENCODER
text_encoder = RNN_ENCODER(algo.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage) #TRAIN.NET_E path can be given directly
text_encoder.load_state_dict(state_dict)
# print('Load text encoder from:', cfg.TRAIN.NET_E) ###edited here
if cfg.CUDA:
text_encoder = text_encoder.cuda()
text_encoder.eval()
#LOADING Generator
netG = G_NET()
model_dir = cfg.TRAIN.NET_G #directory for model can be given directly as well
state_dict = torch.load(model_dir, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
# print('Load G from: ', model_dir) ###edited here
if cfg.CUDA:
netG.cuda()
netG.eval()
return [algo,text_encoder,netG,dataset]
dataset = LSUNClass('%s/%s_%s_lmdb' %
(cfg.DATA_DIR, cfg.DATASET_NAME, split_dir),
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
elif cfg.DATA_DIR.find('imagenet') != -1:
from datasets import ImageFolder
dataset = ImageFolder(cfg.DATA_DIR,
split_dir='train',
custom_classes=CLASS_DIC[cfg.DATASET_NAME],
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
elif cfg.GAN.B_CONDITION: # text to image task
if cfg.DATASET_NAME == 'birds':
from datasets import TextDataset
dataset = TextDataset(cfg.DATA_DIR,
split_dir,
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
elif cfg.DATASET_NAME == 'flowers':
from datasets import FlowersDataset
dataset = FlowersDataset(cfg.DATA_DIR,
split_dir,
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
assert dataset
num_gpu = len(cfg.GPU_ID.split(','))
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=cfg.TRAIN.BATCH_SIZE * num_gpu,
drop_last=True,
shuffle=bshuffle,
num_workers=int(cfg.WORKERS)) #, collate_fn=dataset.collator)
num_imgs = 4
g_dir = cfg.TRAIN.NET_G
data_dir = os.path.join(cfg.DATA_DIR, cfg.DATA_SIZE)
imsize = cfg.TREE.BASE_SIZE * (2**(cfg.TREE.BRANCH_NUM - 1))
image_transform = transforms.Compose([
transforms.Resize(int(imsize * 76 / 64)),
transforms.RandomCrop(imsize),
transforms.RandomHorizontalFlip()
])
data_dir = os.path.join(cfg.DATA_DIR, cfg.DATA_SIZE)
dataset = TextDataset(data_dir,
'test',
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
experimenter_new = Experimenter(embedding_dim=cfg.TEXT.EMBEDDING_DIM,
net_E=cfg.TRAIN.NET_E,
n_words=dataset.n_words,
wordtoix=dataset.wordtoix)
experimenter_original = Experimenter(embedding_dim=cfg.TEXT.EMBEDDING_DIM,
net_E=cfg.TRAIN.NET_E,
n_words=dataset.n_words,
wordtoix=dataset.wordtoix)
sentences = {
'descriptive': [
'a skier with a red jacket on going down the side of a mountain',
if __name__ == "__main__":
args = parse_args()
if args.cfg_file is not None:
cfg_from_file(args.cfg_file)
if args.gpu_id != -1:
cfg.GPU_ID = args.gpu_id
print('Using config:')
pprint.pprint(cfg)
now = datetime.datetime.now(dateutil.tz.tzlocal())
timestamp = now.strftime('%Y_%m_%d_%H_%M_%S')
datadir = 'Data/%s' % cfg.DATASET_NAME
dataset = TextDataset(datadir, cfg.EMBEDDING_TYPE, 1)
filename_test = '%s/test' % (datadir)
dataset.test = dataset.get_data(filename_test)
if cfg.TRAIN.FLAG:
filename_train = '%s/train' % (datadir)
dataset.train = dataset.get_data(filename_train)
ckt_logs_dir = "ckt_logs/%s/%s_%s" % (cfg.DATASET_NAME,
cfg.CONFIG_NAME, timestamp)
mkdir_p(ckt_logs_dir)
else:
s_tmp = cfg.TRAIN.PRETRAINED_MODEL
ckt_logs_dir = s_tmp[:s_tmp.find('.ckpt')]
model = CondGAN(image_shape=dataset.image_shape)
import torch
import torch.nn as nn
import torch.optim as optim
from datasets import TextDataset
from model import LSTMAE
from generate import generate
import numpy as np
data_root = 'Data'
max_length = 50
batch_size = 50
num_epochs = 100
learning_rate = 0.0003
print_interval = 2
dataset = TextDataset(data_root, max_length)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
vocab_size = len(dataset.get_charset())
losses = deque([], maxlen=print_interval)
input_size = max_length
hidden_size = 16
model = LSTMAE(vocab_size, input_size, hidden_size, num_layers=1, isCuda=False)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
loss_func = nn.CrossEntropyLoss()
losses = deque([], maxlen=print_interval)
for epoch in range(num_epochs):
batch_index = 0
for batch_i, samples in enumerate(dataloader):
imsize = cfg.TREE.BASE_SIZE * (2**(cfg.TREE.BRANCH_NUM - 1))
batch_size = cfg.TRAIN.BATCH_SIZE
if cfg.TRAIN.TRANS == 'org':
image_transform = transforms.Compose([
transforms.Scale(int(imsize * 76 / 64)),
transforms.RandomCrop(imsize),
transforms.RandomHorizontalFlip()
])
else:
image_transform = transforms.Compose([
transforms.Scale(int(imsize * 76 / 64)),
transforms.RandomCrop(imsize)
])
dataset = TextDataset(cfg.DATA_DIR,
'train',
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
print(dataset.n_words, dataset.embeddings_num)
assert dataset
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
drop_last=True,
shuffle=True,
num_workers=int(cfg.WORKERS))
# # validation data #
dataset_val = TextDataset(cfg.DATA_DIR,
'test',
base_size=cfg.TREE.BASE_SIZE,
transform=transforms.Compose(
train=cfg.TRAIN.FLAG,
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform,
feature_switch=args.audio_switch)
elif cfg.DATA_DIR.find("birds") != -1:
from datasets import BirdsDataset
dataset = BirdsDataset(cfg.DATA_DIR,
train=cfg.TRAIN.FLAG,
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform,
feature_switch=args.audio_switch)
elif cfg.GAN.B_CONDITION: # text to image task
from datasets import TextDataset
dataset = TextDataset(cfg.DATA_DIR,
split_dir,
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform,
asr_flag=args.asr_flag,
audio_switch=args.audio_switch)
assert dataset
num_gpu = len(cfg.GPU_ID.split(','))
param = {
"num_workers": min(int(cfg.WORKERS), os.cpu_count()),
"pin_memory": True
} if cfg.CUDA else {}
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
dataset)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=cfg.TRAIN.BATCH_SIZE,
##########################################################################
torch.cuda.set_device(cfg.GPU_ID)
cudnn.benchmark = True
# Get data loader ##################################################
imsize = cfg.TREE.BASE_SIZE
batch_size = cfg.TRAIN.BATCH_SIZE
image_transform = transforms.Compose([
transforms.Resize(int(imsize * 76 / 64)),
transforms.RandomCrop(imsize),
transforms.RandomHorizontalFlip()
])
if cfg.B_VALIDATION:
dataset = TextDataset(cfg.DATA_DIR,
'test',
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
logger.info(dataset.n_words, dataset.embeddings_num)
assert dataset
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
drop_last=True,
shuffle=True,
num_workers=int(cfg.WORKERS))
else:
dataset = TextDataset(cfg.DATA_DIR,
'train',
base_size=cfg.TREE.BASE_SIZE,
transform=image_transform)
logger.info(dataset.n_words, dataset.embeddings_num)
assert dataset
if not os.path.exists(args.training_image):
os.makedirs('{:s}'.format(args.training_image))
writer = SummaryWriter()
size = (args.image_size, args.image_size)
train_tf = transforms.Compose([
transforms.Resize(size),
transforms.RandomHorizontalFlip(),
transforms.RandomGrayscale(),
])
dataset_train = TextDataset(args.root,
'train',
base_size=args.base_size,
CAPTIONS_PER_IMAGE=args.CAPTIONS_PER_IMAGE,
WORDS_NUM=args.WORDS_NUM,
BRANCH_NUM=args.BRANCH_NUM,
transform=train_tf)
assert dataset_train
train_set = torch.utils.data.DataLoader(dataset_train,
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
num_workers=args.n_threads)
print(len(train_set))
ixtoword_train = dataset_train.ixtoword
g_model = InpaintNet().to(device)
