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 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
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):
# ################### 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 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 build_models(self):
################### Text and Image encoders ########################################
if cfg.TRAIN.NET_E == '':
print('Error: no pretrained text-image encoders')
return
"""
image_encoder = CNN_dummy()
image_encoder.cuda()
image_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()
"""
VGG = VGG16()
VGG.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 ##############
from model import D_NET256
netD = D_NET256()
netG = EncDecNet()
netD.apply(weights_init)
netG.apply(weights_init)
#
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()
netG.cuda()
netD.cuda()
VGG.cuda()
return [text_encoder, netG, netD, epoch, VGG]
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 = HigherLevelRNN(ninput=cfg.TEXT.EMBEDDING_DIM,
nhidden=cfg.TEXT.EMBEDDING_DIM)
state_dict = torch.load(cfg.TRAIN.NET_E, map_location=lambda s, l: s)
text_encoder.load_state_dict(state_dict)
text_encoder_L = RNN_ENCODER(self.n_words,
nhidden=cfg.TEXT.EMBEDDING_DIM)
L_path = cfg.TRAIN.NET_E.replace('text_encoder', 'text_encoder_L')
state_dict = torch.load(L_path, map_location=lambda s, l: s)
text_encoder_L.load_state_dict(state_dict)
for p in itertools.chain(text_encoder.parameters(),\
text_encoder_L.parameters()):
p.requires_grad = False
print('Loaded text encoder: %s' % cfg.TRAIN.NET_E)
print('Loaded low level text encoder: %s' % L_path)
text_encoder.eval()
text_encoder_L.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)
epoch = state_dict['epoch'] + 1
netG.load_state_dict(state_dict['netG'])
for i in range(len(netsD)):
netsD[i].load_state_dict(state_dict['netD'][i])
# 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()
text_encoder_L = text_encoder_L.cuda()
image_encoder = image_encoder.cuda()
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
return [(text_encoder, text_encoder_L), image_encoder, netG, netsD,
epoch]
def build_models(self):
###########
if cfg.TRAIN.NET_E == '':
print('Error: no pretrained text-image encoders')
return
image_encoder = BERT_CNN_ENCODER_RNN_DECODER(cfg.TEXT.EMBEDDING_DIM,
cfg.CNN_RNN.HIDDEN_DIM,
self.n_words,
rec_unit=cfg.RNN_TYPE)
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 = \
BERT_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()
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]
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()
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):
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)
# ########################################################### #
# Create a target network.
target_netG = deepcopy(netG)
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
netG.cuda()
# The target network is stored on the scondary GPU.---------------------------------
target_netG.cuda()
#target_netG.ca_net.device = secondary_device
#-----------------------------------------------------------------------------------
for i in range(len(netsD)):
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]
def build_models(self):
print('Building models...')
print('N_words: ', self.n_words)
#####################
## 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)
print('Built image encoder: ', image_encoder)
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)
print('Built text encoder: ', text_encoder)
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_encoder
if cfg.CAP.USE_ORIGINAL:
caption_cnn = CAPTION_CNN(embed_size=cfg.TEXT.EMBEDDING_DIM)
caption_rnn = CAPTION_RNN(embed_size=cfg.TEXT.EMBEDDING_DIM,
hidden_size=cfg.CAP.HIDDEN_SIZE,
vocab_size=self.n_words,
num_layers=cfg.CAP.NUM_LAYERS)
else:
caption_cnn = Encoder()
caption_rnn = Decoder(idx2word=self.ixtoword)
caption_cnn_checkpoint = torch.load(
cfg.CAP.CAPTION_CNN_PATH,
map_location=lambda storage, loc: storage)
caption_rnn_checkpoint = torch.load(
cfg.CAP.CAPTION_RNN_PATH,
map_location=lambda storage, loc: storage)
caption_cnn.load_state_dict(caption_cnn_checkpoint['model_state_dict'])
caption_rnn.load_state_dict(caption_rnn_checkpoint['model_state_dict'])
for p in caption_cnn.parameters():
p.requires_grad = False
print('Load caption model from: ', cfg.CAP.CAPTION_CNN_PATH)
caption_cnn.eval()
for p in caption_rnn.parameters():
p.requires_grad = False
print('Load caption model from: ', cfg.CAP.CAPTION_RNN_PATH)
#################################
## GENERATOR & 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]
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)
text_encoder = text_encoder.to(cfg.DEVICE)
image_encoder = image_encoder.to(cfg.DEVICE)
caption_cnn = caption_cnn.to(cfg.DEVICE)
caption_rnn = caption_rnn.to(cfg.DEVICE)
netG.to(cfg.DEVICE)
for i in range(len(netsD)):
netsD[i].to(cfg.DEVICE)
return [
text_encoder, image_encoder, caption_cnn, caption_rnn, netG, netsD,
epoch
]
def load_network(gpus:list, distributed:bool):
netG = G_NET()
netG.apply(weights_init)
if distributed:
netG = netG.cuda()
netG = torch.nn.parallel.DistributedDataParallel(netG, device_ids=gpus, output_device=gpus[0], broadcast_buffers=True)
else:
if cfg.CUDA:
netG = netG.cuda()
netG = torch.nn.DataParallel(netG, device_ids=gpus)
print(netG)
netsD = []
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())
if cfg.TREE.BRANCH_NUM > 3:
netsD.append(D_NET512())
if cfg.TREE.BRANCH_NUM > 4:
netsD.append(D_NET1024())
# TODO: if cfg.TREE.BRANCH_NUM > 5:
# netsD_module = nn.ModuleList(netsD)
# netsD_module.apply(weights_init)
# netsD_module = torch.nn.parallel.DistributedDataParallel(netsD_module.cuda(), device_ids=gpus, output_device=gpus[0])
for i in range(len(netsD)):
netsD[i].apply(weights_init)
if distributed:
netsD[i] = torch.nn.parallel.DistributedDataParallel(netsD[i].cuda(), device_ids=gpus, output_device=gpus[0], broadcast_buffers=True
# , process_group=pg_Ds[i]
)
else:
netsD[i] = torch.nn.DataParallel(netsD[i], device_ids=gpus)
print(netsD[i])
print('# of netsD', len(netsD))
count = 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 ', cfg.TRAIN.NET_G)
istart = cfg.TRAIN.NET_G.rfind('_') + 1
iend = cfg.TRAIN.NET_G.rfind('.')
count = cfg.TRAIN.NET_G[istart:iend]
count = int(count) + 1
if cfg.TRAIN.NET_D != '':
for i in range(len(netsD)):
print('Load %s_%d.pth' % (cfg.TRAIN.NET_D, i))
state_dict = torch.load('%s%d.pth' % (cfg.TRAIN.NET_D, i), map_location=lambda storage, loc: storage)
netsD[i].load_state_dict(state_dict)
inception_model = INCEPTION_V3()
if not distributed:
if cfg.CUDA:
netG.cuda()
for i in range(len(netsD)):
netsD[i].cuda()
inception_model = inception_model.cuda()
inception_model = torch.nn.DataParallel(inception_model, device_ids=gpus)
else:
inception_model = torch.nn.parallel.DistributedDataParallel(inception_model.cuda(), device_ids=gpus, output_device=gpus[0])
pass
# inception_model = inception_model.cpu() #to(torch.device("cuda:{}".format(gpus[0])))
inception_model.eval()
print("model device, G:{}, D:{}, incep:{}".format(netG.device_ids, netsD[0].device_ids, inception_model.device_ids))
return netG, netsD, len(netsD), inception_model, count
