def build_models(self):
# ###################encoders######################################## #
image_encoder = ImageEncoder(output_channels=cfg.hidden_dim)
if cfg.text_encoder_path != '':
img_encoder_path = cfg.text_encoder_path.replace('text_encoder', 'image_encoder')
print('Load image encoder from:', img_encoder_path)
state_dict = torch.load(img_encoder_path, map_location='cpu')
if 'model' in state_dict.keys():
image_encoder.load_state_dict(state_dict['model'])
else:
image_encoder.load_state_dict(state_dict)
for p in image_encoder.parameters(): # make image encoder grad on
p.requires_grad = True
# image_encoder.eval()
epoch = 0
###################################################################
text_encoder = TextEncoder(bert_config = self.bert_config)
if cfg.text_encoder_path != '':
epoch = cfg.text_encoder_path[istart:iend]
epoch = int(epoch) + 1
text_encoder_path = cfg.text_encoder_path
print('Load text encoder from:', text_encoder_path)
state_dict = torch.load(text_encoder_path, map_location='cpu')
if 'model' in state_dict.keys():
text_encoder.load_state_dict(state_dict['model'])
else:
text_encoder.load_state_dict(state_dict)
for p in text_encoder.parameters(): # make text encoder grad on
p.requires_grad = True
# ########################################################### #
if cfg.CUDA:
text_encoder = text_encoder.cuda()
image_encoder = image_encoder.cuda()
return [text_encoder, image_encoder, epoch]
def load_network(self):
image_generator = ImageGenerator()
image_generator.apply(weights_init)
disc_image = DiscriminatorImage()
disc_image.apply(weights_init)
emb_dim = 300
text_encoder = TextEncoder(emb_dim, self.txt_emb,
1, dropout=0.0)
attn_model = 'general'
text_generator = TextGenerator(attn_model, emb_dim, len(self.txt_dico.id2word),
self.txt_emb,
n_layers=1, dropout=0.0)
image_encoder = ImageEncoder()
image_encoder.apply(weights_init)
disc_latent = DiscriminatorLatent(emb_dim)
if cfg.NET_G != '':
state_dict = \
torch.load(cfg.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load from: ', cfg.NET_G)
if cfg.NET_D != '':
state_dict = \
torch.load(cfg.NET_D,
map_location=lambda storage, loc: storage)
netD.load_state_dict(state_dict)
print('Load from: ', cfg.NET_D)
if cfg.ENCODER != '':
state_dict = \
torch.load(cfg.ENCODER,
map_location=lambda storage, loc: storage)
encoder.load_state_dict(state_dict)
print('Load from: ', cfg.ENCODER)
if cfg.DECODER != '':
state_dict = \
torch.load(cfg.DECODER,
map_location=lambda storage, loc: storage)
decoder.load_state_dict(state_dict)
print('Load from: ', cfg.DECODER)
if cfg.IMAGE_ENCODER != '':
state_dict = \
torch.load(cfg.IMAGE_ENCODER,
map_location=lambda storage, loc: storage)
image_encoder.load_state_dict(state_dict)
print('Load from: ', cfg.IMAGE_ENCODER)
if cfg.CUDA:
image_encoder.cuda()
image_generator.cuda()
text_encoder.cuda()
text_generator.cuda()
disc_image.cuda()
disc_latent.cuda()
return image_encoder, image_generator, text_encoder, text_generator, disc_image, disc_latent