def build_models():
# build model ############################################################
text_encoder = RNN_ENCODER(dataset.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
labels = Variable(torch.LongTensor(range(batch_size)))
start_epoch = 0
if cfg.TRAIN.NET_E != '':
state_dict = torch.load(cfg.TRAIN.NET_E)
text_encoder.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_E)
#
name = cfg.TRAIN.NET_E.replace('text_encoder', 'image_encoder')
state_dict = torch.load(name)
image_encoder.load_state_dict(state_dict)
print('Load ', name)
istart = cfg.TRAIN.NET_E.rfind('_') + 8
iend = cfg.TRAIN.NET_E.rfind('.')
start_epoch = cfg.TRAIN.NET_E[istart:iend]
start_epoch = int(start_epoch) + 1
print('start_epoch', start_epoch)
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
labels = labels.cuda()
return text_encoder, image_encoder, labels, start_epoch
def build_models():
text_encoder = RNN_ENCODER(dataset.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
labels = Variable(torch.LongTensor(range(batch_size)))
start_epoch = 0
lr = cfg.TRAIN.ENCODER_LR
if cfg.TRAIN.NET_E != '':
state_dict = torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
print('Load {}'.format(cfg.TRAIN.NET_E))
name = cfg.TRAIN.NET_E.replace('text_encoder', 'image_encoder')
state_dict = torch.load(name, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
print('Load {}'.format(name))
istart = cfg.TRAIN.NET_E.rfind('_') + 8
iend = cfg.TRAIN.NET_E.rfind('.')
start_epoch = cfg.TRAIN.NET_E[istart:iend]
start_epoch = int(start_epoch) + 1
print('start_epoch', start_epoch)
# initial lr with the right value
# note that the turning point is always epoch 114
if start_epoch < 114:
lr = cfg.TRAIN.ENCODER_LR * (0.98 ** start_epoch)
else:
lr = cfg.TRAIN.ENCODER_LR / 10
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
labels = labels.cuda()
return text_encoder, image_encoder, labels, start_epoch, lr
def build_models(text_encoder_type):
# build model ############################################################
text_encoder_type = text_encoder_type.casefold()
if text_encoder_type not in ('rnn', 'transformer'):
raise ValueError('Unsupported text_encoder_type')
if text_encoder_type == 'rnn':
text_encoder = RNN_ENCODER(dataset.n_words,
nhidden=cfg.TEXT.EMBEDDING_DIM)
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
labels = Variable(torch.LongTensor(range(batch_size)))
start_epoch = 0
if cfg.TRAIN.NET_E:
if text_encoder_type == 'rnn':
state_dict = torch.load(cfg.TRAIN.NET_E)
text_encoder.load_state_dict(state_dict)
elif text_encoder_type == 'transformer':
text_encoder = GPT2Model.from_pretrained(cfg.TRAIN.NET_E)
# output_hidden_states = True )
print('Load ', cfg.TRAIN.NET_E)
#
name = cfg.TRAIN.NET_E.replace('text_encoder', 'image_encoder')
state_dict = torch.load(name)
image_encoder.load_state_dict(state_dict)
print('Load ', name)
istart = cfg.TRAIN.NET_E.rfind('_') + 8
iend = cfg.TRAIN.NET_E.rfind('.')
start_epoch = cfg.TRAIN.NET_E[istart:iend]
start_epoch = int(start_epoch) + 1
else:
if text_encoder_type == 'rnn':
print('Training RNN from scratch')
elif text_encoder_type == 'transformer':
# don't initialize the weights of these huge models from scratch...
print('Training Transformer starting from pretrained model')
text_encoder = GPT2Model.from_pretrained(TRANSFORMER_ENCODER)
# output_hidden_states = True )
print('Training CNN starting from ImageNet pretrained Inception-v3')
print('start_epoch', start_epoch)
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
labels = labels.cuda()
return text_encoder, image_encoder, labels, start_epoch
def build_models():
Inception_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
Inception_encoder.init_trainable_weights()
labels = Variable(torch.LongTensor(range(batch_size)))
if cfg.TRAIN.NET_I != '':
Iname = cfg.TRAIN.NET_I
state_dict = torch.load(Iname)
Inception_encoder.load_state_dict(state_dict)
print('Load ', Iname)
if cfg.CUDA:
Inception_encoder = Inception_encoder.cuda()
Inception_encoder = Inception_encoder.eval()
labels = labels.cuda()
return Inception_encoder, labels
def build_models():
# build model ############################################################
# text_encoder = RNN_ENCODER(dataset.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
text_encoder = TEXT_TRANSFORMER_ENCODERv2(emb=cfg.TEXT.EMBEDDING_DIM,
heads=8,
depth=1,
seq_length=cfg.TEXT.WORDS_NUM,
num_tokens=dataset.n_words)
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
labels = Variable(torch.LongTensor(range(batch_size)))
start_epoch = 0
if cfg.TRAIN.NET_E != '':
print('Loading... ', cfg.TRAIN.NET_E)
state_dict = torch.load(cfg.TRAIN.NET_E, map_location='cpu')
# new_state_dict = {}
# for k,v in state_dict.items():
# new_state_dict[k.replace('module.','')]=v
# new_state_dict
text_encoder.load_state_dict(state_dict['model'])
# text_encoder.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_E)
#
name = cfg.TRAIN.NET_E.replace('text_encoder', 'image_encoder')
state_dict = torch.load(name, map_location='cpu')
# new_state_dict = {}
# for k,v in state_dict.items():
# new_state_dict[k.replace('module.','')]=v
# new_state_dict
image_encoder.load_state_dict(state_dict['model'])
print('Load ', name)
istart = cfg.TRAIN.NET_E.rfind('_') + 8
iend = cfg.TRAIN.NET_E.rfind('.')
start_epoch = cfg.TRAIN.NET_E[istart:iend]
start_epoch = int(start_epoch) + 1
print('start_epoch', start_epoch)
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
labels = labels.cuda()
return text_encoder, image_encoder, labels, start_epoch
def build_models():
# build model ############################################################
text_encoder = RNN_ENCODER(dataset.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
labels = Variable(torch.LongTensor(range(batch_size)))
start_epoch = 0
# MODIFIED
if cfg.PRETRAINED_RNN:
text_encoder_params = torch.load(
cfg.PRETRAINED_RNN, map_location=lambda storage, loc: storage)
text_encoder.rnn.load_state_dict(text_encoder_params['encoder'])
pad_idx = text_encoder_params['vocab']['word2id']['<pad>']
n_words, embed_size = text_encoder.encoder.weight.size()
text_encoder.encoder = nn.Embedding(n_words, embed_size, pad_idx)
text_encoder.encoder.load_state_dict(text_encoder_params['embedding'])
if cfg.PRETRAINED_CNN:
image_encoder_params = torch.load(
cfg.PRETRAINED_CNN, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(image_encoder_params)
if cfg.TRAIN.NET_E != '':
state_dict = torch.load(cfg.TRAIN.NET_E)
text_encoder.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_E)
#
name = cfg.TRAIN.NET_E.replace('text_encoder', 'image_encoder')
state_dict = torch.load(name)
image_encoder.load_state_dict(state_dict)
print('Load ', name)
istart = cfg.TRAIN.NET_E.rfind('_') + 8
iend = cfg.TRAIN.NET_E.rfind('.')
start_epoch = cfg.TRAIN.NET_E[istart:iend]
start_epoch = int(start_epoch) + 1
print('start_epoch', start_epoch)
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
labels = labels.cuda()
return text_encoder, image_encoder, labels, start_epoch
def build_models(dataset, batch_size, audio_flag=False):
# build model ############################################################
if audio_flag:
text_encoder = CNNRNN_Attn(40,
nhidden=cfg.TEXT.EMBEDDING_DIM,
nsent=cfg.TEXT.SENT_EMBEDDING_DIM)
else:
text_encoder = RNN_ENCODER(dataset.n_words,
nhidden=cfg.TEXT.EMBEDDING_DIM,
nsent=cfg.TEXT.SENT_EMBEDDING_DIM)
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM,
condition=cfg.TRAIN.MASK_COND,
condition_channel=0)
labels = torch.LongTensor(range(batch_size))
start_epoch = 0
if cfg.TRAIN.NET_E != '':
state_dict = torch.load(cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
print('Load ', cfg.TRAIN.NET_E)
#
name = cfg.TRAIN.NET_E.replace('text_encoder', 'image_encoder')
state_dict = torch.load(name,
map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
print('Load ', name)
# istart = cfg.TRAIN.NET_E.rfind('encoder')
# iend = cfg.TRAIN.NET_E.rfind('.')
start_epoch = re.match(r'.*_encoder(\d+).*', cfg.TRAIN.NET_E).group(1)
start_epoch = int(start_epoch) + 1
print('start_epoch', start_epoch)
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
labels = labels.cuda()
return text_encoder, image_encoder, labels, start_epoch
def build_models(self):
# ###################encoders######################################## #
if cfg.TRAIN.NET_E == '':
print('Error: no pretrained text-image encoders')
return
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters():
p.requires_grad = False
print('Load image encoder from:', img_encoder_path)
image_encoder.eval()
text_encoder = \
RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
for p in text_encoder.parameters():
p.requires_grad = False
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder.eval()
# #######################generator and discriminators############## #
netG = G_NET(len(self.cats_index_dict))
netsPatD, netsShpD = [], []
if cfg.TREE.BRANCH_NUM > 0:
netsPatD.append(PAT_D_NET64())
netsShpD.append(SHP_D_NET64(len(self.cats_index_dict)))
if cfg.TREE.BRANCH_NUM > 1:
netsPatD.append(PAT_D_NET128())
netsShpD.append(SHP_D_NET128(len(self.cats_index_dict)))
if cfg.TREE.BRANCH_NUM > 2:
netsPatD.append(PAT_D_NET256())
netsShpD.append(SHP_D_NET256(len(self.cats_index_dict)))
netObjSSD = OBJ_SS_D_NET(len(self.cats_index_dict))
netObjLSD = OBJ_LS_D_NET(len(self.cats_index_dict))
netG.apply(weights_init)
netObjSSD.apply(weights_init)
netObjLSD.apply(weights_init)
for i in range(len(netsPatD)):
netsPatD[i].apply(weights_init)
netsShpD[i].apply(weights_init)
print('# of netsPatD', len(netsPatD))
# ########################################################### #
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
netG.cuda()
netObjSSD.cuda()
netObjLSD.cuda()
for i in range(len(netsPatD)):
netsPatD[i].cuda()
netsShpD[i].cuda()
if len(cfg.GPU_IDS) > 1:
text_encoder = nn.DataParallel(text_encoder)
text_encoder.to(self.device)
image_encoder = nn.DataParallel(image_encoder)
image_encoder.to(self.device)
netG = nn.DataParallel(netG)
netG.to(self.device)
netObjSSD = nn.DataParallel(netObjSSD)
netObjSSD.to(self.device)
netObjLSD = nn.DataParallel(netObjLSD)
netObjLSD.to(self.device)
for i in range(len(netsPatD)):
netsPatD[i] = nn.DataParallel(netsPatD[i])
netsPatD[i].to(self.device)
netsShpD[i] = nn.DataParallel(netsShpD[i])
netsShpD[i].to(self.device)
#
epoch = 0
if cfg.TRAIN.NET_G != '':
state_dict = \
torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
epoch = cfg.TRAIN.NET_G[istart:iend]
epoch = int(epoch) + 1
Gname = cfg.TRAIN.NET_G
for i in range(len(netsPatD)):
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netPatD%d.pth' % (s_tmp, i)
print('Load PatD from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netsPatD[i].load_state_dict(state_dict)
Dname = '%s/netShpD%d.pth' % (s_tmp, i)
print('Load ShpD from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netsShpD[i].load_state_dict(state_dict)
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netObjSSD.pth' % (s_tmp)
print('Load ObjSSD from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netObjSSD.load_state_dict(state_dict)
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netObjLSD.pth' % (s_tmp)
print('Load ObjLSD from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netObjLSD.load_state_dict(state_dict)
return [
text_encoder, image_encoder, netG, netsPatD, netsShpD, netObjSSD,
netObjLSD, epoch
]
def build_models(self):
# text encoders
if cfg.TRAIN.NET_E == '':
print('Error: no pretrained text-image encoders')
return
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder', 'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters():
p.requires_grad = False
print('Load image encoder from:', img_encoder_path)
image_encoder.eval()
# self.n_words = 156
text_encoder = RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
for p in text_encoder.parameters():
p.requires_grad = False
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder.eval()
# Caption models - cnn_encoder and rnn_decoder
caption_cnn = CAPTION_CNN(cfg.CAP.embed_size)
caption_cnn.load_state_dict(torch.load(cfg.CAP.caption_cnn_path, map_location=lambda storage, loc: storage))
for p in caption_cnn.parameters():
p.requires_grad = False
print('Load caption model from:', cfg.CAP.caption_cnn_path)
caption_cnn.eval()
# self.n_words = 9
caption_rnn = CAPTION_RNN(cfg.CAP.embed_size, cfg.CAP.hidden_size * 2, self.n_words, cfg.CAP.num_layers)
# caption_rnn = CAPTION_RNN(cfg.CAP.embed_size, cfg.CAP.hidden_size * 2, self.n_words, cfg.CAP.num_layers)
caption_rnn.load_state_dict(torch.load(cfg.CAP.caption_rnn_path, map_location=lambda storage, loc: storage))
for p in caption_rnn.parameters():
p.requires_grad = False
print('Load caption model from:', cfg.CAP.caption_rnn_path)
# Generator and Discriminator:
netsD = []
if cfg.GAN.B_DCGAN:
if cfg.TREE.BRANCH_NUM == 1:
from model import D_NET64 as D_NET
elif cfg.TREE.BRANCH_NUM == 2:
from model import D_NET128 as D_NET
else: # cfg.TREE.BRANCH_NUM == 3:
from model import D_NET256 as D_NET
netG = G_DCGAN()
netsD = [D_NET(b_jcu=False)]
else:
from model import D_NET64, D_NET128, D_NET256
netG = G_NET()
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
netG.apply(weights_init)
# print(netG)
for i in range(len(netsD)):
netsD[i].apply(weights_init)
# print(netsD[i])
print('# of netsD', len(netsD))
epoch = 0
if cfg.TRAIN.NET_G != '':
state_dict = \
torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
epoch = cfg.TRAIN.NET_G[istart:iend]
# print(epoch)
# print(state_dict.keys())
# print(netG.keys())
# epoch = state_dict['epoch']
epoch = int(epoch) + 1
# epoch = 187
if cfg.TRAIN.B_NET_D:
Gname = cfg.TRAIN.NET_G
for i in range(len(netsD)):
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netD%d.pth' % (s_tmp, i)
print('Load D from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netsD[i].load_state_dict(state_dict)
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
caption_cnn = caption_cnn.cuda()
caption_rnn = caption_rnn.cuda()
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
return [text_encoder, image_encoder, caption_cnn, caption_rnn, netG, netsD, epoch]
def build_models(self):
# ###################encoders######################################## #
if cfg.TRAIN.NET_E == '':
raise FileNotFoundError(
'No pretrained text encoder found in directory DAMSMencoders/. \n'
+
'Please train the DAMSM first before training the GAN (see README for details).'
)
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters():
p.requires_grad = False
print('Load image encoder from:', img_encoder_path)
image_encoder.eval()
if self.text_encoder_type == 'rnn':
text_encoder = \
RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
elif self.text_encoder_type == 'transformer':
text_encoder = GPT2Model.from_pretrained(TRANSFORMER_ENCODER)
state_dict = \
torch.load(cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
for p in text_encoder.parameters():
p.requires_grad = False
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder.eval()
# #######################generator and discriminators############## #
netsD = []
if cfg.GAN.B_DCGAN:
if cfg.TREE.BRANCH_NUM == 1:
from model import D_NET64 as D_NET
elif cfg.TREE.BRANCH_NUM == 2:
from model import D_NET128 as D_NET
else: # cfg.TREE.BRANCH_NUM == 3:
from model import D_NET256 as D_NET
# TODO: elif cfg.TREE.BRANCH_NUM > 3:
netG = G_DCGAN()
netsD = [D_NET(b_jcu=False)]
elif cfg.GAN.B_STYLEGEN:
netG = G_NET_STYLED()
if cfg.GAN.B_STYLEDISC:
from model import D_NET_STYLED64, D_NET_STYLED128, D_NET_STYLED256
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET_STYLED64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET_STYLED128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET_STYLED256())
# TODO: if cfg.TREE.BRANCH_NUM > 3:
else:
from model import D_NET64, D_NET128, D_NET256
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
# TODO: if cfg.TREE.BRANCH_NUM > 3:
else:
from model import D_NET64, D_NET128, D_NET256
netG = G_NET()
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
# TODO: if cfg.TREE.BRANCH_NUM > 3:
netG.apply(weights_init)
# print(netG)
for i in range(len(netsD)):
netsD[i].apply(weights_init)
# print(netsD[i])
print(netG.__class__)
for i in netsD:
print(i.__class__)
print('# of netsD', len(netsD))
#
epoch = 0
if cfg.TRAIN.NET_G != '':
state_dict = torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
if cfg.GAN.B_STYLEGEN:
netG.w_ewma = state_dict['w_ewma']
if cfg.CUDA:
netG.w_ewma = netG.w_ewma.to('cuda:' + str(cfg.GPU_ID))
netG.load_state_dict(state_dict['netG_state_dict'])
else:
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
epoch = cfg.TRAIN.NET_G[istart:iend]
epoch = int(epoch) + 1
if cfg.TRAIN.B_NET_D:
Gname = cfg.TRAIN.NET_G
for i in range(len(netsD)):
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netD%d.pth' % (s_tmp, i)
print('Load D from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netsD[i].load_state_dict(state_dict)
# ########################################################### #
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
return [text_encoder, image_encoder, netG, netsD, epoch]
def calc_mp(self):
if cfg.TRAIN.NET_G == '':
print('Error: the path for main module is not found!')
else:
#if split_dir == 'test':
# split_dir = 'valid'
if cfg.GAN.B_DCGAN:
netG = G_DCGAN()
else:
netG = EncDecNet()
netG.apply(weights_init)
netG.cuda()
netG.eval()
# The text encoder
text_encoder = RNN_ENCODER(self.n_words,
nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder = text_encoder.cuda()
text_encoder.eval()
# The image encoder
#image_encoder = CNN_dummy()
#print('define image_encoder')
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
print('Load image encoder from:', img_encoder_path)
image_encoder = image_encoder.cuda()
image_encoder.eval()
# The VGG network
VGG = VGG16()
print("Load the VGG model")
#VGG.to(torch.device("cuda:1"))
VGG.cuda()
VGG.eval()
batch_size = self.batch_size
nz = cfg.GAN.Z_DIM
noise = Variable(torch.FloatTensor(batch_size, nz), volatile=True)
noise = noise.cuda()
model_dir = os.path.join(cfg.DATA_DIR, 'output', self.args.netG,
'Model', self.args.netG_epoch)
state_dict = \
torch.load(model_dir, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', model_dir)
# the path to save modified images
cnt = 0
idx = 0
diffs, sims = [], []
for _ in range(1): # (cfg.TEXT.CAPTIONS_PER_IMAGE):
for step, data in enumerate(self.data_loader, 0):
cnt += batch_size
if step % 100 == 0:
print('step: ', step)
imgs, w_imgs, captions, cap_lens, class_ids, keys, wrong_caps, \
wrong_caps_len, wrong_cls_id = prepare_data(data)
#######################################################
# (1) Extract text and image embeddings
######################################################
hidden = text_encoder.init_hidden(batch_size)
words_embs, sent_emb = text_encoder(
wrong_caps, wrong_caps_len, hidden)
words_embs, sent_emb = words_embs.detach(
), sent_emb.detach()
mask = (wrong_caps == 0)
num_words = words_embs.size(2)
if mask.size(1) > num_words:
mask = mask[:, :num_words]
#######################################################
# (2) Modify real images
######################################################
noise.data.normal_(0, 1)
fake_img, mu, logvar = netG(imgs[-1], sent_emb, words_embs,
noise, mask, VGG)
diff = F.l1_loss(fake_img, imgs[-1])
diffs.append(diff.item())
region_code, cnn_code = image_encoder(fake_img)
sim = cosine_similarity(sent_emb, cnn_code)
sim = torch.mean(sim)
sims.append(sim.item())
diff = np.sum(diffs) / len(diffs)
sim = np.sum(sims) / len(sims)
print('diff: %.3f, sim:%.3f' % (diff, sim))
print('MP: %.3f' % ((1 - diff) * sim))
netG_epoch = self.args.netG_epoch[self.args.netG_epoch.find('_') +
1:-4]
print('model_epoch:%s, diff: %.3f, sim:%.3f, MP:%.3f' %
(netG_epoch, np.sum(diffs) / len(diffs),
np.sum(sims) / len(sims), (1 - diff) * sim))
def build_models(self):
################### Text and Image encoders ########################################
# if cfg.TRAIN.NET_E == '':
# print('Error: no pretrained text-image encoders')
# return
VGG = VGGNet()
for p in VGG.parameters():
p.requires_grad = False
print("Load the VGG model")
VGG.eval()
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
text_encoder = RNN_ENCODER(self.n_words,
nhidden=cfg.TEXT.EMBEDDING_DIM)
# when NET_E = '', train the image_encoder and text_encoder jointly
if cfg.TRAIN.NET_E != '':
state_dict = torch.load(
cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage).state_dict()
text_encoder.load_state_dict(state_dict)
for p in text_encoder.parameters():
p.requires_grad = False
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder.eval()
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = torch.load(
img_encoder_path,
map_location=lambda storage, loc: storage).state_dict()
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters():
p.requires_grad = False
print('Load image encoder from:', img_encoder_path)
image_encoder.eval()
####################### Generator and Discriminators ##############
netsD = []
if cfg.GAN.B_DCGAN:
if cfg.TREE.BRANCH_NUM == 1:
from model import D_NET64 as D_NET
elif cfg.TREE.BRANCH_NUM == 2:
from model import D_NET128 as D_NET
else: # cfg.TREE.BRANCH_NUM == 3:
from model import D_NET256 as D_NET
netG = G_DCGAN()
if cfg.TRAIN.W_GAN:
netsD = [D_NET(b_jcu=False)]
else:
from model import D_NET64, D_NET128, D_NET256
netG = G_NET()
netG.apply(weights_init)
if cfg.TRAIN.W_GAN:
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
for i in range(len(netsD)):
netsD[i].apply(weights_init)
print('# of netsD', len(netsD))
#
epoch = 0
if cfg.TRAIN.NET_G != '':
state_dict = torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
epoch = cfg.TRAIN.NET_G[istart:iend]
epoch = int(epoch) + 1
if cfg.TRAIN.B_NET_D:
Gname = cfg.TRAIN.NET_G
for i in range(len(netsD)):
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netD%d.pth' % (s_tmp, i)
print('Load D from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netsD[i].load_state_dict(state_dict)
# ########################################################### #
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
netG.cuda()
VGG = VGG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
return [text_encoder, image_encoder, netG, netsD, epoch, VGG]
def gen_example(self, data_dic):
if cfg.TRAIN.NET_G == '' or cfg.TRAIN.NET_C == '':
print('Error: the path for main module or DCM is not found!')
else:
# The text encoder
text_encoder = \
RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder = text_encoder.cuda()
text_encoder.eval()
# The image encoder
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
print('Load image encoder from:', img_encoder_path)
image_encoder = image_encoder.cuda()
image_encoder.eval()
# The VGG network
VGG = VGGNet()
print("Load the VGG model")
VGG.cuda()
VGG.eval()
# The main module
if cfg.GAN.B_DCGAN:
netG = G_DCGAN()
else:
netG = G_NET()
s_tmp = cfg.TRAIN.NET_G[:cfg.TRAIN.NET_G.rfind('.pth')]
model_dir = cfg.TRAIN.NET_G
state_dict = \
torch.load(model_dir, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', model_dir)
netG.cuda()
netG.eval()
# The DCM
netDCM = DCM_Net()
if cfg.TRAIN.NET_C != '':
state_dict = \
torch.load(cfg.TRAIN.NET_C, map_location=lambda storage, loc: storage)
netDCM.load_state_dict(state_dict)
print('Load DCM from: ', cfg.TRAIN.NET_C)
netDCM.cuda()
netDCM.eval()
for key in data_dic:
save_dir = '%s/%s' % (s_tmp, key)
mkdir_p(save_dir)
captions, cap_lens, sorted_indices, imgs = data_dic[key]
batch_size = captions.shape[0]
nz = cfg.GAN.Z_DIM
captions = Variable(torch.from_numpy(captions), volatile=True)
cap_lens = Variable(torch.from_numpy(cap_lens), volatile=True)
captions = captions.cuda()
cap_lens = cap_lens.cuda()
for i in range(1):
noise = Variable(torch.FloatTensor(batch_size, nz),
volatile=True)
noise = noise.cuda()
#######################################################
# (1) Extract text and image embeddings
######################################################
hidden = text_encoder.init_hidden(batch_size)
# The text embeddings
words_embs, sent_emb = text_encoder(
captions, cap_lens, hidden)
# The image embeddings
region_features, cnn_code = \
image_encoder(imgs[cfg.TREE.BRANCH_NUM - 1].unsqueeze(0))
mask = (captions == 0)
#######################################################
# (2) Modify real images
######################################################
noise.data.normal_(0, 1)
fake_imgs, attention_maps, mu, logvar, h_code, c_code = netG(
noise, sent_emb, words_embs, mask, cnn_code,
region_features)
real_img = imgs[cfg.TREE.BRANCH_NUM - 1].unsqueeze(0)
real_features = VGG(real_img)[0]
fake_img = netDCM(h_code, real_features, sent_emb, words_embs, \
mask, c_code)
cap_lens_np = cap_lens.cpu().data.numpy()
for j in range(batch_size):
save_name = '%s/%d_s_%d' % (save_dir, i,
sorted_indices[j])
for k in range(len(fake_imgs)):
im = fake_imgs[k][j].data.cpu().numpy()
im = (im + 1.0) * 127.5
im = im.astype(np.uint8)
im = np.transpose(im, (1, 2, 0))
im = Image.fromarray(im)
fullpath = '%s_g%d.png' % (save_name, k)
im.save(fullpath)
for k in range(len(attention_maps)):
if len(fake_imgs) > 1:
im = fake_imgs[k + 1].detach().cpu()
else:
im = fake_imgs[0].detach().cpu()
attn_maps = attention_maps[k]
att_sze = attn_maps.size(2)
img_set, sentences = \
build_super_images2(im[j].unsqueeze(0),
captions[j].unsqueeze(0),
[cap_lens_np[j]], self.ixtoword,
[attn_maps[j]], att_sze)
if img_set is not None:
im = Image.fromarray(img_set)
fullpath = '%s_a%d.png' % (save_name, k)
im.save(fullpath)
save_name = '%s/%d_sf_%d' % (save_dir, 1,
sorted_indices[j])
im = fake_img[j].data.cpu().numpy()
im = (im + 1.0) * 127.5
im = im.astype(np.uint8)
im = np.transpose(im, (1, 2, 0))
im = Image.fromarray(im)
fullpath = '%s_SF.png' % (save_name)
im.save(fullpath)
save_name = '%s/%d_s_%d' % (save_dir, 1, 9)
im = imgs[2].data.cpu().numpy()
im = (im + 1.0) * 127.5
im = im.astype(np.uint8)
im = np.transpose(im, (1, 2, 0))
im = Image.fromarray(im)
fullpath = '%s_SR.png' % (save_name)
im.save(fullpath)
def build_models(self):
# ###################encoders######################################## #
if cfg.TRAIN.NET_E == '':
print('Error: no pretrained text-image encoders')
return
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters():
p.requires_grad = False
print('Load image encoder from:', img_encoder_path)
image_encoder.eval()
text_encoder = \
RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
for p in text_encoder.parameters():
p.requires_grad = False
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder.eval()
# #######################generator and discriminators############## #
netsD = []
if cfg.GAN.B_DCGAN:
if cfg.TREE.BRANCH_NUM == 1:
from model import D_NET64 as D_NET
elif cfg.TREE.BRANCH_NUM == 2:
from model import D_NET128 as D_NET
else: # cfg.TREE.BRANCH_NUM == 3:
from model import D_NET256 as D_NET
# TODO: elif cfg.TREE.BRANCH_NUM > 3:
netG = G_DCGAN()
netsD = [D_NET(b_jcu=False)]
else:
from model import D_NET64, D_NET128, D_NET256
netG = G_NET()
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
# TODO: if cfg.TREE.BRANCH_NUM > 3:
netG.apply(weights_init)
# print(netG)
for i in range(len(netsD)):
netsD[i].apply(weights_init)
# print(netsD[i])
print('# of netsD', len(netsD))
#
epoch = 0
# MODIFIED
if cfg.PRETRAINED_G != '':
state_dict = torch.load(cfg.PRETRAINED_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.PRETRAINED_G)
if cfg.TRAIN.B_NET_D:
Gname = cfg.PRETRAINED_G
s_tmp = Gname[:Gname.rfind('/')]
for i in range(len(netsD)):
Dname = '%s/netD%d.pth' % (
s_tmp, i
) # the name of Ds should be consistent and differ from each other in i
print('Load D from: ', Dname)
state_dict = torch.load(
Dname, map_location=lambda storage, loc: storage)
netsD[i].load_state_dict(state_dict)
if cfg.TRAIN.NET_G != '':
state_dict = \
torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
epoch = cfg.TRAIN.NET_G[istart:iend]
epoch = int(epoch) + 1
if cfg.TRAIN.B_NET_D:
Gname = cfg.TRAIN.NET_G
for i in range(len(netsD)):
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netD%d.pth' % (s_tmp, i)
print('Load D from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netsD[i].load_state_dict(state_dict)
# ########################################################### #
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
return [text_encoder, image_encoder, netG, netsD, epoch]
def build_models(self):
# ###################encoders######################################## #
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
image_encoder.train()
# #######################generator and discriminators############## #
netsD = []
if cfg.GAN.B_DCGAN:
if cfg.TREE.BRANCH_NUM == 1:
from model import D_NET64 as D_NET
elif cfg.TREE.BRANCH_NUM == 2:
from model import D_NET128 as D_NET
else: # cfg.TREE.BRANCH_NUM == 3:
from model import D_NET256 as D_NET
# TODO: elif cfg.TREE.BRANCH_NUM > 3:
netG = G_DCGAN()
netsD = [D_NET(b_jcu=False)]
else:
from model import D_NET64, D_NET128, D_NET256
netG = G_NET()
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
# TODO: if cfg.TREE.BRANCH_NUM > 3:
netG.apply(weights_init)
# print(netG)
for i in range(len(netsD)):
netsD[i].apply(weights_init)
# print(netsD[i])
print('# of netsD', len(netsD))
#
epoch = 0
if cfg.PRETRAINED_CNN:
image_encoder_params = torch.load(
cfg.PRETRAINED_CNN, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(image_encoder_params)
if cfg.PRETRAINED_G != '':
state_dict = torch.load(cfg.PRETRAINED_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.PRETRAINED_G)
if cfg.TRAIN.B_NET_D:
Gname = cfg.PRETRAINED_G
s_tmp = Gname[:Gname.rfind('/')]
for i in range(len(netsD)):
Dname = '%s/netD%d.pth' % (
s_tmp, i
) # the name of Ds should be consistent and differ from each other in i
print('Load D from: ', Dname)
state_dict = torch.load(
Dname, map_location=lambda storage, loc: storage)
netsD[i].load_state_dict(state_dict)
# ########################################################### #
if cfg.CUDA:
image_encoder = image_encoder.cuda()
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
return [image_encoder, netG, netsD, epoch]
def sampling(self, split_dir):
if cfg.TRAIN.NET_G == '':
print('Error: the path for morels is not found!')
else:
if split_dir == 'test':
split_dir = 'valid'
# Build and load the generator
if cfg.GAN.B_DCGAN:
netG = G_DCGAN()
else:
netG = G_NET()
netG.apply(weights_init)
netG.cuda()
netG.eval()
# load text encoder
text_encoder = RNN_ENCODER(self.n_words,
nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = torch.load(cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder = text_encoder.cuda()
text_encoder.eval()
#load image encoder
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = torch.load(img_encoder_path,
map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
print('Load image encoder from:', img_encoder_path)
image_encoder = image_encoder.cuda()
image_encoder.eval()
batch_size = self.batch_size
nz = cfg.GAN.Z_DIM
noise = Variable(torch.FloatTensor(batch_size, nz), volatile=True)
noise = noise.cuda()
model_dir = cfg.TRAIN.NET_G
state_dict = torch.load(model_dir,
map_location=lambda storage, loc: storage)
# state_dict = torch.load(cfg.TRAIN.NET_G)
netG.load_state_dict(state_dict)
print('Load G from: ', model_dir)
# the path to save generated images
s_tmp = model_dir[:model_dir.rfind('.pth')]
save_dir = '%s/%s' % (s_tmp, split_dir)
mkdir_p(save_dir)
cnt = 0
R_count = 0
R = np.zeros(30000)
cont = True
for ii in range(11): # (cfg.TEXT.CAPTIONS_PER_IMAGE):
if (cont == False):
break
for step, data in enumerate(self.data_loader, 0):
cnt += batch_size
if (cont == False):
break
if step % 100 == 0:
print('cnt: ', cnt)
# if step > 50:
# break
imgs, captions, cap_lens, class_ids, keys = prepare_data(
data)
hidden = text_encoder.init_hidden(batch_size)
# words_embs: batch_size x nef x seq_len
# sent_emb: batch_size x nef
words_embs, sent_emb = text_encoder(
captions, cap_lens, hidden)
words_embs, sent_emb = words_embs.detach(
), sent_emb.detach()
mask = (captions == 0)
num_words = words_embs.size(2)
if mask.size(1) > num_words:
mask = mask[:, :num_words]
#######################################################
# (2) Generate fake images
######################################################
noise.data.normal_(0, 1)
fake_imgs, _, _, _ = netG(noise, sent_emb, words_embs,
mask, cap_lens)
for j in range(batch_size):
s_tmp = '%s/single/%s' % (save_dir, keys[j])
folder = s_tmp[:s_tmp.rfind('/')]
if not os.path.isdir(folder):
#print('Make a new folder: ', folder)
mkdir_p(folder)
k = -1
# for k in range(len(fake_imgs)):
im = fake_imgs[k][j].data.cpu().numpy()
# [-1, 1] --> [0, 255]
im = (im + 1.0) * 127.5
im = im.astype(np.uint8)
im = np.transpose(im, (1, 2, 0))
im = Image.fromarray(im)
fullpath = '%s_s%d_%d.png' % (s_tmp, k, ii)
im.save(fullpath)
_, cnn_code = image_encoder(fake_imgs[-1])
for i in range(batch_size):
mis_captions, mis_captions_len = self.dataset.get_mis_caption(
class_ids[i])
hidden = text_encoder.init_hidden(99)
_, sent_emb_t = text_encoder(mis_captions,
mis_captions_len, hidden)
rnn_code = torch.cat(
(sent_emb[i, :].unsqueeze(0), sent_emb_t), 0)
### cnn_code = 1 * nef
### rnn_code = 100 * nef
scores = torch.mm(cnn_code[i].unsqueeze(0),
rnn_code.transpose(0, 1)) # 1* 100
cnn_code_norm = torch.norm(cnn_code[i].unsqueeze(0),
2,
dim=1,
keepdim=True)
rnn_code_norm = torch.norm(rnn_code,
2,
dim=1,
keepdim=True)
norm = torch.mm(cnn_code_norm,
rnn_code_norm.transpose(0, 1))
scores0 = scores / norm.clamp(min=1e-8)
if torch.argmax(scores0) == 0:
R[R_count] = 1
R_count += 1
if R_count >= 30000:
sum = np.zeros(10)
np.random.shuffle(R)
for i in range(10):
sum[i] = np.average(R[i * 3000:(i + 1) * 3000 - 1])
R_mean = np.average(sum)
R_std = np.std(sum)
print("R mean:{:.4f} std:{:.4f}".format(R_mean, R_std))
cont = False
def embedding(self, split_dir, model):
if cfg.TRAIN.NET_G == '':
print('Error: the path for morels is not found!')
else:
if split_dir == 'test':
split_dir = 'valid'
# Build and load the generator
if cfg.GAN.B_DCGAN:
netG = G_DCGAN()
else:
netG = G_NET()
netG.apply(weights_init)
if cfg.GPU_ID != -1:
netG.cuda()
netG.eval()
#
model_dir = cfg.TRAIN.NET_G
state_dict = \
torch.load(model_dir, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', model_dir)
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
print(img_encoder_path)
print('Load image encoder from:', img_encoder_path)
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
if cfg.GPU_ID != -1:
image_encoder = image_encoder.cuda()
image_encoder.eval()
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder = RNN_ENCODER(self.n_words,
nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
if cfg.GPU_ID != -1:
text_encoder = text_encoder.cuda()
text_encoder.eval()
batch_size = self.batch_size
nz = cfg.GAN.Z_DIM
with torch.no_grad():
noise = Variable(torch.FloatTensor(batch_size, nz))
if cfg.GPU_ID != -1:
noise = noise.cuda()
# the path to save generated images
save_dir = model_dir[:model_dir.rfind('.pth')]
cnt = 0
# new
if cfg.TRAIN.CLIP_SENTENCODER:
print("Use CLIP SentEncoder for sampling")
img_features = dict()
txt_features = dict()
with torch.no_grad():
for _ in range(1): # (cfg.TEXT.CAPTIONS_PER_IMAGE):
for step, data in enumerate(self.data_loader, 0):
cnt += batch_size
if step % 100 == 0:
print('step: ', step)
imgs, captions, cap_lens, class_ids, keys, texts = prepare_data(
data)
hidden = text_encoder.init_hidden(batch_size)
# words_embs: batch_size x nef x seq_len
# sent_emb: batch_size x nef
words_embs, sent_emb = text_encoder(
captions, cap_lens, hidden)
words_embs, sent_emb = words_embs.detach(
), sent_emb.detach()
mask = (captions == 0)
num_words = words_embs.size(2)
if mask.size(1) > num_words:
mask = mask[:, :num_words]
if cfg.TRAIN.CLIP_SENTENCODER:
# random select one paragraph for each training example
sents = []
for idx in range(len(texts)):
sents_per_image = texts[idx].split(
'\n') # new 3/11
if len(sents_per_image) > 1:
sent_ix = np.random.randint(
0,
len(sents_per_image) - 1)
else:
sent_ix = 0
sents.append(sents_per_image[0])
# print('sents: ', sents)
sent = clip.tokenize(sents) # .to(device)
# load clip
#model = torch.jit.load("model.pt").cuda().eval()
sent_input = sent
if cfg.GPU_ID != -1:
sent_input = sent.cuda()
# print("text input", sent_input)
sent_emb_clip = model.encode_text(
sent_input).float()
if CLIP:
sent_emb = sent_emb_clip
#######################################################
# (2) Generate fake images
######################################################
noise.data.normal_(0, 1)
fake_imgs, _, _, _ = netG(noise, sent_emb, words_embs,
mask)
if CLIP:
images = []
for j in range(fake_imgs[-1].shape[0]):
image = fake_imgs[-1][j].cpu().clone()
image = image.squeeze(0)
unloader = transforms.ToPILImage()
image = unloader(image)
image = preprocess(
image.convert("RGB")) # 256*256 -> 224*224
images.append(image)
image_mean = torch.tensor(
[0.48145466, 0.4578275, 0.40821073]).cuda()
image_std = torch.tensor(
[0.26862954, 0.26130258, 0.27577711]).cuda()
image_input = torch.tensor(np.stack(images)).cuda()
image_input -= image_mean[:, None, None]
image_input /= image_std[:, None, None]
cnn_codes = model.encode_image(image_input).float()
else:
region_features, cnn_codes = image_encoder(
fake_imgs[-1])
for j in range(batch_size):
cnn_code = cnn_codes[j]
temp = keys[j].replace('b', '').replace("'", '')
img_features[temp] = cnn_code.cpu().numpy()
txt_features[temp] = sent_emb[j].cpu().numpy()
with open(save_dir + ".pkl", 'wb') as f:
pickle.dump(img_features, f)
with open(save_dir + "_text.pkl", 'wb') as f:
pickle.dump(txt_features, f)
def build_models(self):
# ###################encoders######################################## #
if cfg.TRAIN.NET_E == '':
print('Error: no pretrained text-image encoders')
return
####################################################################
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters(): # make image encoder grad on
p.requires_grad = True
for k, v in image_encoder.named_children(
): # freeze the layer1-5 (set eval for BNlayer)
if k in frozen_list_image_encoder:
v.train(False)
v.requires_grad_(False)
print('Load image encoder from:', img_encoder_path)
# image_encoder.eval()
###################################################################
text_encoder = TEXT_TRANSFORMER_ENCODERv2(
emb=cfg.TEXT.EMBEDDING_DIM,
heads=8,
depth=1,
seq_length=cfg.TEXT.WORDS_NUM,
num_tokens=self.n_words)
# state_dict = torch.load(cfg.TRAIN.NET_E)
# text_encoder.load_state_dict(state_dict)
# print('Load ', cfg.TRAIN.NET_E)
state_dict = torch.load(cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
for p in text_encoder.parameters():
p.requires_grad = True
print('Load text encoder from:', cfg.TRAIN.NET_E)
# text_encoder.eval()
# #######################generator and discriminators############## #
netsD = []
if cfg.GAN.B_DCGAN:
if cfg.TREE.BRANCH_NUM == 1:
from model import D_NET64 as D_NET
elif cfg.TREE.BRANCH_NUM == 2:
from model import D_NET128 as D_NET
else: # cfg.TREE.BRANCH_NUM == 3:
from model import D_NET256 as D_NET
# TODO: elif cfg.TREE.BRANCH_NUM > 3:
netG = G_DCGAN()
netsD = [D_NET(b_jcu=False)]
else:
from model import D_NET64, D_NET128, D_NET256
netG = G_NET()
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
# TODO: if cfg.TREE.BRANCH_NUM > 3:
netG.apply(weights_init)
# print(netG)
for i in range(len(netsD)):
netsD[i].apply(weights_init)
# print(netsD[i])
print('# of netsD', len(netsD))
#
epoch = 0
if cfg.TRAIN.NET_G != '':
state_dict = \
torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
epoch = cfg.TRAIN.NET_G[istart:iend]
epoch = int(epoch) + 1
if cfg.TRAIN.B_NET_D:
Gname = cfg.TRAIN.NET_G
for i in range(len(netsD)):
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netD%d.pth' % (s_tmp, i)
print('Load D from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netsD[i].load_state_dict(state_dict)
# ########################################################## #
# config = Config()
cap_model = caption.build_model_v3(config)
print("Initializing from Checkpoint...")
cap_model_path = cfg.TRAIN.NET_E.replace('text_encoder', 'cap_model')
if os.path.exists(cap_model_path):
print('Load C from: {0}'.format(cap_model_path))
state_dict = \
torch.load(cap_model_path, map_location=lambda storage, loc: storage)
cap_model.load_state_dict(state_dict['model'])
else:
base_line_path = 'catr/checkpoints/catr_damsm256_proj_coco2014_ep02.pth'
print('Load C from: {0}'.format(base_line_path))
checkv3 = torch.load(base_line_path, map_location='cpu')
cap_model.load_state_dict(checkv3['model'], strict=False)
# ########################################################### #
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
cap_model = cap_model.cuda() # caption transformer added
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
return [text_encoder, image_encoder, netG, netsD, epoch, cap_model]
def build_models(self):
# ################### models ######################################## #
if cfg.TRAIN.NET_E == '':
print('Error: no pretrained text-image encoders')
return
if cfg.TRAIN.NET_G == '':
print('Error: no pretrained main module')
return
VGG = VGGNet()
for p in VGG.parameters():
p.requires_grad = False
print("Load the VGG model")
VGG.eval()
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters():
p.requires_grad = False
print('Load image encoder from:', img_encoder_path)
image_encoder.eval()
text_encoder = \
RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
for p in text_encoder.parameters():
p.requires_grad = False
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder.eval()
if cfg.GAN.B_DCGAN:
netG = G_DCGAN()
from model import D_NET256 as D_NET
netD = D_NET(b_jcu=False)
else:
from model import D_NET256
netG = G_NET()
netD = D_NET256()
netD.apply(weights_init)
state_dict = \
torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
netG.eval()
print('Load G from: ', cfg.TRAIN.NET_G)
epoch = 0
netDCM = DCM_Net()
if cfg.TRAIN.NET_C != '':
state_dict = \
torch.load(cfg.TRAIN.NET_C, map_location=lambda storage, loc: storage)
netDCM.load_state_dict(state_dict)
print('Load DCM from: ', cfg.TRAIN.NET_C)
istart = cfg.TRAIN.NET_C.rfind('_') + 1
iend = cfg.TRAIN.NET_C.rfind('.')
epoch = cfg.TRAIN.NET_C[istart:iend]
epoch = int(epoch) + 1
if cfg.TRAIN.NET_D != '':
state_dict = \
torch.load(cfg.TRAIN.NET_D, map_location=lambda storage, loc: storage)
netD.load_state_dict(state_dict)
print('Load DCM Discriminator from: ', cfg.TRAIN.NET_D)
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
netG.cuda()
netDCM.cuda()
VGG = VGG.cuda()
netD.cuda()
return [text_encoder, image_encoder, netG, netD, epoch, VGG, netDCM]
def sampling(self, split_dir):
if cfg.TRAIN.NET_G == '': #TODO
print('Error: the path for morels is not found!')
else:
if split_dir == 'test':
split_dir = 'valid'
# Build and load the generator
if cfg.GAN.B_DCGAN:
netG = G_DCGAN()
else:
netG = G_NET()
netG.apply(weights_init)
netG.cuda()
netG.eval()
#
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder') #TODO
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
print('Load image encoder from:', img_encoder_path)
image_encoder = image_encoder.cuda()
image_encoder.eval()
batch_size = self.batch_size
nz = cfg.GAN.Z_DIM
noise = Variable(torch.FloatTensor(batch_size, nz), volatile=True)
noise = noise.cuda()
model_dir = cfg.TRAIN.NET_G
state_dict = \
torch.load(model_dir, map_location=lambda storage, loc: storage)
# state_dict = torch.load(cfg.TRAIN.NET_G)
netG.load_state_dict(state_dict)
print('Load G from: ', model_dir)
# the path to save generated images
s_tmp = model_dir[:model_dir.rfind('.pth')]
save_dir = '%s/%s' % (s_tmp, split_dir)
mkdir_p(save_dir)
cnt = 0
words_embs = Variable(
torch.zeros(batch_size, cfg.TEXT.EMBEDDING_DIM,
cfg.TEXT.WORDS_NUM))
mask = Variable(torch.zeros(batch_size, cfg.TEXT.WORDS_NUM))
words_embs, mask = words_embs.cuda(), mask.cuda()
for _ in range(1): # (cfg.TEXT.CAPTIONS_PER_IMAGE):
for step, data in enumerate(self.data_loader, 0):
cnt += batch_size
if step % 100 == 0:
print('step: ', step)
# if step > 50:
# break
imgs, keys = prepare_data(data)
# words_embs: batch_size x nef x seq_len
# sent_emb: batch_size x nef
_, sent_emb = image_encoder(imgs[-1])
sent_emb = sent_emb.detach()
#######################################################
# (2) Generate fake images
######################################################
noise.data.normal_(0, 1)
fake_imgs, _, _, _ = netG(noise, sent_emb, words_embs,
mask)
for j in range(batch_size):
s_tmp = '%s/single/%s' % (save_dir, keys[j])
folder = s_tmp[:s_tmp.rfind('/')]
if not os.path.isdir(folder):
print('Make a new folder: ', folder)
mkdir_p(folder)
k = -1
# for k in range(len(fake_imgs)):
im = fake_imgs[k][j].data.cpu().numpy()
# [-1, 1] --> [0, 255]
im = (im + 1.0) * 127.5
im = im.astype(np.uint8)
im = np.transpose(im, (1, 2, 0))
im = Image.fromarray(im)
fullpath = '%s_s%d.png' % (s_tmp, k)
im.save(fullpath)
def build_models(self):
# ###################encoders######################################## #
if cfg.TRAIN.NET_E == '':
self.logger.error('Error: no pretrained text-image encoders')
return
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM,
condition=cfg.TRAIN.MASK_COND,
condition_channel=1)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters():
p.requires_grad = False
self.logger.info(
'Load image encoder from: {}'.format(img_encoder_path))
image_encoder.eval()
if self.audio_flag:
text_encoder = CNNRNN_Attn(n_filters=40,
nhidden=cfg.TEXT.EMBEDDING_DIM,
nsent=cfg.TEXT.SENT_EMBEDDING_DIM)
else:
text_encoder = \
RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM, nsent=cfg.TEXT.SENT_EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
for p in text_encoder.parameters():
p.requires_grad = False
self.logger.info('Load text encoder from: {}'.format(cfg.TRAIN.NET_E))
text_encoder.eval()
# #######################generator and discriminators############## #
netsD = []
if cfg.GAN.B_DCGAN:
if cfg.TREE.BRANCH_NUM == 1:
from model import D_NET64 as D_NET
elif cfg.TREE.BRANCH_NUM == 2:
from model import D_NET128 as D_NET
else: # cfg.TREE.BRANCH_NUM == 3:
from model import D_NET256 as D_NET
# TODO: elif cfg.TREE.BRANCH_NUM > 3:
netG = G_DCGAN()
netsD = [D_NET(b_jcu=False)]
else:
from model import D_NET64, D_NET128, D_NET256
netG = G_NET()
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
# TODO: if cfg.TREE.BRANCH_NUM > 3:
netG.apply(weights_init)
# print(netG)
for i in range(len(netsD)):
netsD[i].apply(weights_init)
# print(netsD[i])
self.logger.info('# of netsD: {}'.format(len(netsD)))
#
epoch = 0
if cfg.TRAIN.NET_G != '':
state_dict = \
torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
self.logger.info('Load G from: {}'.format(cfg.TRAIN.NET_G))
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
epoch = cfg.TRAIN.NET_G[istart:iend]
epoch = int(epoch) + 1
if cfg.TRAIN.B_NET_D:
Gname = cfg.TRAIN.NET_G
for i in range(len(netsD)):
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netD%d.pth' % (s_tmp, i)
self.logger.info('Load D from: {}'.format(Dname))
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netsD[i].load_state_dict(state_dict)
# ########################################################### #
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
return [text_encoder, image_encoder, netG, netsD, epoch]
def sampling(self, split_dir):
if cfg.TRAIN.NET_G == '':
print('Error: the path for morels is not found!')
else:
if split_dir == 'test':
split_dir = 'valid'
# Build and load the generator
if cfg.GAN.B_DCGAN:
netG = G_DCGAN()
else:
netG = G_NET()
netG.apply(weights_init)
netG.cuda()
netG.eval()
# load text encoder
text_encoder = RNN_ENCODER(self.n_words,
nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = torch.load(cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder = text_encoder.cuda()
text_encoder.eval()
#load image encoder
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = torch.load(img_encoder_path,
map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
print('Load image encoder from:', img_encoder_path)
image_encoder = image_encoder.cuda()
image_encoder.eval()
batch_size = self.batch_size
nz = cfg.GAN.Z_DIM
noise = Variable(torch.FloatTensor(batch_size, nz), volatile=True)
noise = noise.cuda()
model_dir = cfg.TRAIN.NET_G
state_dict = torch.load(model_dir,
map_location=lambda storage, loc: storage)
# state_dict = torch.load(cfg.TRAIN.NET_G)
netG.load_state_dict(state_dict)
print('Load G from: ', model_dir)
# the path to save generated images
s_tmp = model_dir[:model_dir.rfind('.pth')]
save_dir = '%s/%s' % (s_tmp, split_dir)
mkdir_p(save_dir)
cnt = 0
R_count = 0
R = np.zeros(30000)
cont = True
for ii in range(11): # (cfg.TEXT.CAPTIONS_PER_IMAGE):
if (cont == False):
break
for step, data in enumerate(self.data_loader, 0):
cnt += batch_size
if (cont == False):
break
if step % 100 == 0:
print('cnt: ', cnt)
# if step > 50:
# break
imgs, captions, cap_lens, class_ids, keys = prepare_data(
data)
hidden = text_encoder.init_hidden(batch_size)
# words_embs: batch_size x nef x seq_len
# sent_emb: batch_size x nef
words_embs, sent_emb = text_encoder(
captions, cap_lens, hidden)
words_embs, sent_emb = words_embs.detach(
), sent_emb.detach()
mask = (captions == 0)
num_words = words_embs.size(2)
if mask.size(1) > num_words:
mask = mask[:, :num_words]
#######################################################
# (2) Generate fake images
######################################################
noise.data.normal_(0, 1)
fake_imgs, _, _, _ = netG(noise, sent_emb, words_embs,
mask, cap_lens)
for j in range(batch_size):
s_tmp = '%s/single/%s' % (save_dir, keys[j])
folder = s_tmp[:s_tmp.rfind('/')]
if not os.path.isdir(folder):
#print('Make a new folder: ', folder)
mkdir_p(folder)
k = -1
# for k in range(len(fake_imgs)):
im = fake_imgs[k][j].data.cpu().numpy()
# [-1, 1] --> [0, 255]
im = (im + 1.0) * 127.5
im = im.astype(np.uint8)
im = np.transpose(im, (1, 2, 0))
im = Image.fromarray(im)
fullpath = '%s_s%d_%d.png' % (s_tmp, k, ii)
im.save(fullpath)
if cnt >= 30000:
cont = False
print('Load image encoder from:', img_encoder_path)
image_encoder.eval()
text_encoder = RNN_ENCODER(dataset_train.n_words, nhidden=args.image_size)
text_encoder_path = '../DAMSMencoders/' + args.dataset + '/text_encoder200.pth'
state_dict = torch.load(text_encoder_path,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
for p in text_encoder.parameters():
p.requires_grad = False
print('Load text encoder from:', text_encoder_path)
text_encoder.eval()
if use_cuda:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
def prepare_labels():
batch_size = args.batch_size
real_labels = Variable(torch.FloatTensor(batch_size).fill_(1))
fake_labels = Variable(torch.FloatTensor(batch_size).fill_(0))
match_labels = Variable(torch.LongTensor(range(batch_size)))
if use_cuda:
real_labels = real_labels.cuda()
fake_labels = fake_labels.cuda()
match_labels = match_labels.cuda()
return real_labels, fake_labels, match_labels
def gen_example(self, data_dic):
if cfg.TRAIN.NET_G == '':
print('Error: the path for models is not found!')
else:
# Build and load the generator
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
print('Load image encoder from:', img_encoder_path)
image_encoder = image_encoder.cuda()
image_encoder.eval()
# the path to save generated images
if cfg.GAN.B_DCGAN:
netG = G_DCGAN()
else:
netG = G_NET()
s_tmp = cfg.TRAIN.NET_G[:cfg.TRAIN.NET_G.rfind('.pth')]
model_dir = cfg.TRAIN.NET_G
state_dict = \
torch.load(model_dir, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', model_dir)
netG.cuda()
netG.eval()
words_embs = Variable(
torch.zeros(1, cfg.TEXT.EMBEDDING_DIM, cfg.TEXT.WORDS_NUM))
mask = Variable(torch.zeros(1, cfg.TEXT.WORDS_NUM))
noise = Variable(torch.FloatTensor(1, cfg.GAN.Z_DIM),
volatile=True)
words_embs, mask, noise = words_embs.cuda(), mask.cuda(
), noise.cuda()
for key in data_dic:
save_path = '%s/custom/%s' % (s_tmp, key)
img = data_dic[key]
img = Variable(img).unsqueeze(0).cuda()
#######################################################
# (1) Extract image embeddings
######################################################
_, sent_emb = image_encoder(img)
#######################################################
# (2) Generate fake images
######################################################
noise.data.normal_(0, 1)
fake_imgs, attention_maps, _, _ = netG(noise, sent_emb,
words_embs, mask)
# G attention
for k in range(len(fake_imgs)):
im = fake_imgs[k][-1].data.cpu().numpy()
im = (im + 1.0) * 127.5
im = im.astype(np.uint8)
# print('im', im.shape)
im = np.transpose(im, (1, 2, 0))
# print('im', im.shape)
im = Image.fromarray(im)
fullpath = '%s_g%d.png' % (save_path, k)
im.save(fullpath)
def build_models(self):
def count_parameters(model):
total_param = 0
for name, param in model.named_parameters():
if param.requires_grad:
num_param = np.prod(param.size())
if param.dim() > 1:
print(name, ':',
'x'.join(str(x) for x in list(param.size())),
'=', num_param)
else:
print(name, ':', num_param)
total_param += num_param
return total_param
# ###################encoders######################################## #
if cfg.TRAIN.NET_E == '':
print('Error: no pretrained text-image encoders')
return
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters():
p.requires_grad = False
print('Load image encoder from:', img_encoder_path)
image_encoder.eval()
text_encoder = \
RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
for p in text_encoder.parameters():
p.requires_grad = False
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder.eval()
# #######################generator and discriminators############## #
netsD = []
if cfg.GAN.B_DCGAN:
if cfg.TREE.BRANCH_NUM == 1:
from model import D_NET64 as D_NET
elif cfg.TREE.BRANCH_NUM == 2:
from model import D_NET128 as D_NET
else: # cfg.TREE.BRANCH_NUM == 3:
from model import D_NET256 as D_NET
# TODO: elif cfg.TREE.BRANCH_NUM > 3:
netG = G_DCGAN()
netsD = [D_NET(b_jcu=False)]
else:
from model import D_NET64, D_NET128, D_NET256
netG = G_NET()
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
# TODO: if cfg.TREE.BRANCH_NUM > 3:
print('number of trainable parameters =', count_parameters(netG))
print('number of trainable parameters =', count_parameters(netsD[-1]))
netG.apply(weights_init)
# print(netG)
for i in range(len(netsD)):
netsD[i].apply(weights_init)
# print(netsD[i])
print('# of netsD', len(netsD))
#
epoch = 0
if cfg.TRAIN.NET_G != '':
state_dict = \
torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
epoch = cfg.TRAIN.NET_G[istart:iend]
epoch = int(epoch) + 1
if cfg.TRAIN.B_NET_D:
Gname = cfg.TRAIN.NET_G
for i in range(len(netsD)):
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netD%d.pth' % (s_tmp, i)
print('Load D from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netsD[i].load_state_dict(state_dict)
# ########################################################### #
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
return [text_encoder, image_encoder, netG, netsD, epoch]
def build_models(self):
# ############################## encoders ############################# #
text_encoder = RNN_ENCODER(self.n_words,
nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder.eval()
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters():
p.requires_grad = False
print('Load image encoder from:', img_encoder_path)
image_encoder.eval()
# ########### image generator and (potential) shape generator ########## #
netG = G_NET(len(self.cats_index_dict))
netG.apply(weights_init)
netG.eval()
netShpG = None
if cfg.TEST.USE_GT_BOX_SEG > 0:
netShpG = SHP_G_NET(len(self.cats_index_dict))
netShpG.apply(weights_init)
netShpG.eval()
# ################### parallization and initialization ################## #
if cfg.CUDA:
text_encoder.cuda()
image_encoder.cuda()
netG.cuda()
if cfg.TEST.USE_GT_BOX_SEG > 0:
netShpG.cuda()
if len(cfg.GPU_IDS) > 1:
text_encoder = nn.DataParallel(text_encoder)
text_encoder.to(self.device)
image_encoder = nn.DataParallel(image_encoder)
image_encoder.to(self.device)
netG = nn.DataParallel(netG)
netG.to(self.device)
if cfg.TEST.USE_GT_BOX_SEG > 0:
netShpG = nn.DataParallel(netShpG)
netShpG.to(self.device)
state_dict = torch.load(cfg.TRAIN.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.TRAIN.NET_G)
if cfg.TEST.USE_GT_BOX_SEG > 0:
state_dict = torch.load(cfg.TEST.NET_SHP_G,
map_location=lambda storage, loc: storage)
netShpG.load_state_dict(state_dict)
print('Load Shape G from: ', cfg.TEST.NET_SHP_G)
return [text_encoder, image_encoder, netG, netShpG]
def build_models(self):
if cfg.TRAIN.NET_E == '':
print('Error: no pretrained text-image encoders')
return
# vgg16 network
style_loss = VGGNet()
for p in style_loss.parameters():
p.requires_grad = False
print("Load the style loss model")
style_loss.eval()
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters():
p.requires_grad = False
print('Load image encoder from:', img_encoder_path)
image_encoder.eval()
text_encoder = \
RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
for p in text_encoder.parameters():
p.requires_grad = False
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder.eval()
netsD = []
if cfg.GAN.B_DCGAN:
if cfg.TREE.BRANCH_NUM == 1:
from model import D_NET64 as D_NET
elif cfg.TREE.BRANCH_NUM == 2:
from model import D_NET128 as D_NET
else: # cfg.TREE.BRANCH_NUM == 3:
from model import D_NET256 as D_NET
netG = G_DCGAN()
netsD = [D_NET(b_jcu=False)]
else:
from model import D_NET64, D_NET128, D_NET256
netG = G_NET()
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
netG.apply(weights_init)
for i in range(len(netsD)):
netsD[i].apply(weights_init)
print('# of netsD', len(netsD))
#
epoch = 0
if cfg.TRAIN.NET_G != '':
state_dict = \
torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
epoch = cfg.TRAIN.NET_G[istart:iend]
epoch = int(epoch) + 1
if cfg.TRAIN.B_NET_D:
Gname = cfg.TRAIN.NET_G
for i in range(len(netsD)):
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netD%d.pth' % (s_tmp, i)
print('Load D from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netsD[i].load_state_dict(state_dict)
# Create a target network.
target_netG = deepcopy(netG)
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
style_loss = style_loss.cuda()
# The target network is stored on the scondary GPU.---------------------------------
target_netG.cuda(secondary_device)
target_netG.ca_net.device = secondary_device
#-----------------------------------------------------------------------------------
netG.cuda()
for i in range(len(netsD)):
netsD[i] = netsD[i].cuda()
# Disable training in the target network:
for p in target_netG.parameters():
p.requires_grad = False
return [
text_encoder, image_encoder, netG, target_netG, netsD, epoch,
style_loss
]
def build_models(self):
##feature extractor
inception_model = CNN_ENCODER()
#inception_model = CNN_dummy()
## classifier networks
classifiers = []
for i in range(cfg.ATT_NUM):
cls_model = Att_Classifier()
classifiers.append(cls_model)
# #######################generator and discriminators############## #
netsG = []
netsD = []
if cfg.GAN.B_DCGAN:
if cfg.TREE.BRANCH_NUM == 1:
from model import D_NET64 as D_NET
elif cfg.TREE.BRANCH_NUM == 2:
from model import D_NET128 as D_NET
else: # cfg.TREE.BRANCH_NUM == 3:
from model import D_NET256 as D_NET
# TODO: elif cfg.TREE.BRANCH_NUM > 3:
netG = G_DCGAN()
netsD = [D_NET(b_jcu=False)]
else:
from model import D_NET64, D_NET128, D_NET256
"""
if not self.args.kl_loss:
netG = G_NET_not_CA(self.args)
else:
netG = G_NET()
"""
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
netsG.append(G_NET_not_CA_stage1(self.args))
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
netsG.append(G_NET_not_CA_stage2(self.args))
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
netsG.append(G_NET_not_CA_stage3(self.args))
# TODO: if cfg.TREE.BRANCH_NUM > 3:
for i in range(len(netsG)):
netsG[i].apply(weights_init)
# print(netG)
for i in range(len(netsD)):
netsD[i].apply(weights_init)
# print(netsD[i])
print('# of netsD', len(netsD))
#
epoch = 0
if cfg.TRAIN.NET_G != '':
state_dict = \
torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
epoch = cfg.TRAIN.NET_G[istart:iend]
epoch = int(epoch) + 1
if cfg.TRAIN.B_NET_D:
Gname = cfg.TRAIN.NET_G
for i in range(len(netsD)):
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netD%d.pth' % (s_tmp, i)
print('Load D from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netsD[i].load_state_dict(state_dict)
# ########################################################### #
if cfg.CUDA:
inception_model.cuda()
for i in range(len(netsG)):
netsG[i].cuda()
netsD[i].cuda()
for i in range(len(classifiers)):
classifiers[i].cuda()
return [netsG, netsD, inception_model, classifiers, epoch]
def build_models(self):
# ###################encoders######################################## #
if cfg.TRAIN.NET_E == '':
print('Error: no pretrained text-image encoders')
return
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder', 'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters():
p.requires_grad = False
print('Load image encoder from:', img_encoder_path)
image_encoder.eval()
text_encoder = \
RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E,
map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
for p in text_encoder.parameters():
p.requires_grad = False
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder.eval()
# #######################generator and discriminators############## #
netsD = []
from model import D_NET64, D_NET128, D_NET256
netG = G_NET()
if cfg.TREE.BRANCH_NUM > 0:
netsD.append(D_NET64())
if cfg.TREE.BRANCH_NUM > 1:
netsD.append(D_NET128())
if cfg.TREE.BRANCH_NUM > 2:
netsD.append(D_NET256())
netG.apply(weights_init)
# print(netG)
for i in range(len(netsD)):
netsD[i].apply(weights_init)
# print(netsD[i])
print('# of netsD', len(netsD))
epoch = 0
if self.resume:
checkpoint_list = sorted([ckpt for ckpt in glob.glob(self.model_dir + "/" + '*.pth')])
latest_checkpoint = checkpoint_list[-1]
state_dict = torch.load(latest_checkpoint, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict["netG"])
for i in range(len(netsD)):
netsD[i].load_state_dict(state_dict["netD"][i])
epoch = int(latest_checkpoint[-8:-4]) + 1
print("Resuming training from checkpoint {} at epoch {}.".format(latest_checkpoint, epoch))
#
if cfg.TRAIN.NET_G != '':
state_dict = \
torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
epoch = cfg.TRAIN.NET_G[istart:iend]
epoch = int(epoch) + 1
if cfg.TRAIN.B_NET_D:
Gname = cfg.TRAIN.NET_G
for i in range(len(netsD)):
s_tmp = Gname[:Gname.rfind('/')]
Dname = '%s/netD%d.pth' % (s_tmp, i)
print('Load D from: ', Dname)
state_dict = \
torch.load(Dname, map_location=lambda storage, loc: storage)
netsD[i].load_state_dict(state_dict)
# ########################################################### #
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
return [text_encoder, image_encoder, netG, netsD, epoch]
def sampling(self, split_dir):
if cfg.TRAIN.NET_G == '' or cfg.TRAIN.NET_C == '':
print('Error: the path for main module or DCM is not found!')
else:
if split_dir == 'test':
split_dir = 'valid'
if cfg.GAN.B_DCGAN:
netG = G_DCGAN()
else:
netG = G_NET()
netG.apply(weights_init)
netG.cuda()
netG.eval()
# The text encoder
text_encoder = RNN_ENCODER(self.n_words,
nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = \
torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
text_encoder.load_state_dict(state_dict)
print('Load text encoder from:', cfg.TRAIN.NET_E)
text_encoder = text_encoder.cuda()
text_encoder.eval()
# The image encoder
image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
'image_encoder')
state_dict = \
torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
print('Load image encoder from:', img_encoder_path)
image_encoder = image_encoder.cuda()
image_encoder.eval()
# The VGG network
VGG = VGGNet()
print("Load the VGG model")
VGG.cuda()
VGG.eval()
batch_size = self.batch_size
nz = cfg.GAN.Z_DIM
noise = Variable(torch.FloatTensor(batch_size, nz), volatile=True)
noise = noise.cuda()
# The DCM
netDCM = DCM_Net()
if cfg.TRAIN.NET_C != '':
state_dict = \
torch.load(cfg.TRAIN.NET_C, map_location=lambda storage, loc: storage)
netDCM.load_state_dict(state_dict)
print('Load DCM from: ', cfg.TRAIN.NET_C)
netDCM.cuda()
netDCM.eval()
model_dir = cfg.TRAIN.NET_G
state_dict = \
torch.load(model_dir, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load G from: ', model_dir)
# the path to save modified images
s_tmp = model_dir[:model_dir.rfind('.pth')]
save_dir = '%s/%s' % (s_tmp, split_dir)
mkdir_p(save_dir)
cnt = 0
idx = 0
for _ in range(5): # (cfg.TEXT.CAPTIONS_PER_IMAGE):
for step, data in enumerate(self.data_loader, 0):
cnt += batch_size
if step % 100 == 0:
print('step: ', step)
imgs, captions, cap_lens, class_ids, keys, wrong_caps, \
wrong_caps_len, wrong_cls_id = prepare_data(data)
#######################################################
# (1) Extract text and image embeddings
######################################################
hidden = text_encoder.init_hidden(batch_size)
words_embs, sent_emb = text_encoder(
wrong_caps, wrong_caps_len, hidden)
words_embs, sent_emb = words_embs.detach(
), sent_emb.detach()
mask = (wrong_caps == 0)
num_words = words_embs.size(2)
if mask.size(1) > num_words:
mask = mask[:, :num_words]
region_features, cnn_code = \
image_encoder(imgs[cfg.TREE.BRANCH_NUM - 1])
#######################################################
# (2) Modify real images
######################################################
noise.data.normal_(0, 1)
fake_imgs, attention_maps, mu, logvar, h_code, c_code = netG(
noise, sent_emb, words_embs, mask, cnn_code,
region_features)
real_img = imgs[cfg.TREE.BRANCH_NUM - 1]
real_features = VGG(real_img)[0]
fake_img = netDCM(h_code, real_features, sent_emb, words_embs, \
mask, c_code)
for j in range(batch_size):
s_tmp = '%s/single' % (save_dir)
folder = s_tmp[:s_tmp.rfind('/')]
if not os.path.isdir(folder):
print('Make a new folder: ', folder)
mkdir_p(folder)
k = -1
im = fake_img[j].data.cpu().numpy()
im = (im + 1.0) * 127.5
im = im.astype(np.uint8)
im = np.transpose(im, (1, 2, 0))
im = Image.fromarray(im)
fullpath = '%s_s%d.png' % (s_tmp, idx)
idx = idx + 1
im.save(fullpath)
