def load_network_stageI(self):
from model import STAGE1_G, STAGE1_D
netG = STAGE1_G()
netG.apply(weights_init)
print(netG)
netD = STAGE1_D()
netD.apply(weights_init)
print(netD)
if NET_G != '':
state_dict = \
torch.load(NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load from: ', NET_G)
if NET_D != '':
state_dict = \
torch.load(NET_D,
map_location=lambda storage, loc: storage)
netD.load_state_dict(state_dict)
print('Load from: ', NET_D)
if CUDA:
netG.cuda()
netD.cuda()
return netG, netD
def load_network_stageI(self):
from model import STAGE1_G, STAGE1_D
netG = STAGE1_G()
netG.apply(weights_init)
# print(netG)
netD = STAGE1_D(self.new_arch)
netD.apply(weights_init)
# print(netD)
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.CTModel != '':
state_dict=\
torch.load(cfg.CTModel,
map_location=lambda storage, loc: storage)
self.CTallmodel.load_state_dict(state_dict)
print('Load CT Model From: ', cfg.CTModel)
# print(self.CTallmodel)
if cfg.CUDA:
netG.cuda()
netD.cuda()
return netG, netD
def load_network_stageI(self):
from model import STAGE1_G, STAGE1_D
netG = STAGE1_G()
netG.apply(weights_init)
print(netG)
netD = STAGE1_D()
netD.apply(weights_init)
print(netD)
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.CUDA:
netG.cuda()
netD.cuda()
start_epoch = 0
if cfg.NET_G != '':
istart = cfg.NET_G.rfind('_') + 1
iend = cfg.NET_G.rfind('.')
start_epoch = cfg.NET_G[istart:iend]
start_epoch = int(start_epoch) + 1
return netG, netD, start_epoch
def load_network_stageII(self, text_dim, gf_dim, condition_dim, z_dim, df_dim, res_num):
from model import STAGE1_G, STAGE2_G, STAGE2_D
Stage1_G = STAGE1_G(text_dim, gf_dim, condition_dim, z_dim, self.cuda)
netG = STAGE2_G(Stage1_G, text_dim, gf_dim, condition_dim, z_dim, res_num, self.cuda)
netG.apply(weights_init)
print(netG)
if self.net_g != '':
state_dict = torch.load(self.net_g)
#torch.loadself.net_g, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load from: ', self.net_g)
elif self.stage1_g != '':
state_dict = torch.load(self.stage1_g, map_location=lambda storage, loc: storage)
netG.STAGE1_G.load_state_dict(state_dict)
print('Load from: ', self.stage1_g)
else:
print("Please give the Stage1_G path")
return
netD = STAGE2_D(df_dim, condition_dim)
netD.apply(weights_init)
if self.net_d != '':
state_dict = torch.load(self.net_d, map_location=lambda storage, loc: storage)
netD.load_state_dict(state_dict)
print('Load from: ', self.net_d)
print(netD)
if self.cuda:
netG.cuda()
netD.cuda()
return netG, netD
def load_network_stageII(self):
from model import STAGE1_G, STAGE2_G, STAGE2_D
Stage1_G = STAGE1_G()
netG = STAGE2_G(Stage1_G)
netG.apply(weights_init)
print(netG)
if cfg.NET_G != '':
state_dict = torch.load(cfg.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict["netG"])
print('Load from: ', cfg.NET_G)
elif cfg.STAGE1_G != '':
state_dict = torch.load(cfg.STAGE1_G,
map_location=lambda storage, loc: storage)
netG.STAGE1_G.load_state_dict(state_dict["netG"])
print('Load from: ', cfg.STAGE1_G)
else:
print("Please give the Stage1_G path")
return
netD = STAGE2_D()
netD.apply(weights_init)
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)
print(netD)
if cfg.CUDA:
netG.cuda()
netD.cuda()
return netG, netD
def load_model(model_name):
from model import STAGE1_G, STAGE2_G
stage1_g = STAGE1_G()
netG = STAGE2_G(stage1_g)
state_dict = torch.load(model_name,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
return netG
def sample_s1_image(self, condition, noise, reload_state=False):
from .model import STAGE1_G
if self.sample_transfer_Stage1Gen is None:
self.sample_transfer_Stage1Gen = STAGE1_G()
self.sample_transfer_Stage1Gen.eval()
self.sample_transfer_Stage1Gen_state_dict = torch.load(cfg.STAGE1_G, map_location=lambda storage, loc: storage)
self.sample_transfer_Stage1Gen.load_state_dict(self.sample_transfer_Stage1Gen_state_dict)
else:
if reload_state:
self.sample_transfer_Stage1Gen = STAGE1_G()
self.sample_transfer_Stage1Gen.eval()
self.sample_transfer_Stage1Gen.load_state_dict(self.sample_transfer_Stage1Gen_state_dict)
_, img, _, _ = self.sample_transfer_Stage1Gen(torch.Tensor(condition), torch.Tensor(noise))
return img.permute(0, 2, 3, 1).detach().numpy()
def load_network_stageII(self):
from model import STAGE1_G, STAGE2_G, STAGE2_D
Stage1_G = STAGE1_G()
netG = STAGE2_G(Stage1_G)
netG.apply(weights_init)
print(netG)
print(
"---------------------------------------------------------------------------------"
)
print("Current Working Directory : ", os.getcwd())
print("cfg.NET_G : ", cfg.NET_G)
print(
"---------------------------------------------------------------------------------"
)
if cfg.NET_G != '':
print("THIS THE CURRENT DIRECTORY : ", os.getcwd())
state_dict = torch.load(cfg.NET_G,
map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict["netG"])
print('Load from: ', cfg.NET_G)
elif cfg.STAGE1_G != '':
state_dict = torch.load(cfg.STAGE1_G,
map_location=lambda storage, loc: storage)
netG.STAGE1_G.load_state_dict(state_dict["netG"])
print('Load from: ', cfg.STAGE1_G)
else:
print("Please give the Stage1_G path")
return
netD = STAGE2_D()
netD.apply(weights_init)
if cfg.NET_D != '':
state_dict = \
torch.load(cfg.NET_D,
map_location=lambda storage, loc: storage)
netD.load_state_dict(state_dict["netD"])
print('Load from: ', cfg.NET_D)
print(netD)
if cfg.CUDA:
netG.cuda()
netD.cuda()
return netG, netD
def load_network_stageII(self):
from .model import STAGE1_G, STAGE2_G, STAGE2_D
Stage1_G = STAGE1_G()
for param in Stage1_G.parameters():
param.requires_grad = False
netG = STAGE2_G()
netG.apply(weights_init)
print(netG)
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.STAGE1_G != '':
state_dict = \
torch.load(cfg.STAGE1_G,
map_location=lambda storage, loc: storage)
Stage1_G.load_state_dict(state_dict)
print('Load from: ', cfg.STAGE1_G)
else:
print("Please give the Stage1_G path")
return
netD = STAGE2_D()
netD.apply(weights_init)
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)
print(netD)
if cfg.CUDA:
Stage1_G.cuda(3)
netG.cuda()
netD.cuda()
return Stage1_G, netG, netD
def load_network_stageI(self, text_dim, gf_dim, condition_dim, z_dim, df_dim):
from model import STAGE1_G, STAGE1_D
netG = STAGE1_G(text_dim, gf_dim, condition_dim, z_dim, self.cuda)
netG.apply(weights_init)
print(netG)
netD = STAGE1_D(df_dim, condition_dim)
netD.apply(weights_init)
print(netD)
if self.net_g != '':
state_dict = torch.load(self.net_g)
#torch.load(self.net_g, map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)
print('Load from: ', self.net_g)
if self.net_d != '':
state_dict = torch.load(self.net_d, map_location=lambda storage, loc: storage)
netD.load_state_dict(state_dict)
print('Load from: ', self.net_d)
if self.cuda:
netG.cuda()
netD.cuda()
return netG, netD
def load_network_stageI(self):
from model import STAGE1_G, STAGE1_D
from model import EncoderRNN, LuongAttnDecoderRNN
from model import STAGE1_ImageEncoder, EncodingDiscriminator
netG = STAGE1_G()
netG.apply(weights_init)
#print(netG)
netD = STAGE1_D()
netD.apply(weights_init)
#print(netD)
emb_dim = 300
encoder = EncoderRNN(emb_dim, self.txt_emb, 1, dropout=0.0)
attn_model = 'general'
decoder = LuongAttnDecoderRNN(attn_model,
emb_dim,
len(self.txt_dico.id2word),
self.txt_emb,
n_layers=1,
dropout=0.0)
image_encoder = STAGE1_ImageEncoder()
image_encoder.apply(weights_init)
e_disc = EncodingDiscriminator(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)
# load classification model and freeze weights
#clf_model = models.alexnet(pretrained=True)
clf_model = models.vgg16(pretrained=True)
for param in clf_model.parameters():
param.requires_grad = False
if cfg.CUDA:
netG.cuda()
netD.cuda()
encoder.cuda()
decoder.cuda()
image_encoder.cuda()
e_disc.cuda()
clf_model.cuda()
# ## finetune model for a bit???
# output_size = 512 * 2 * 2
# num_classes = 200
# clf_model.classifier = nn.Sequential(
# nn.Linear(output_size, 1024, bias=True),
# nn.LeakyReLU(0.2),
# nn.Dropout(0.5),
# nn.Linear(1024, num_classes, bias=True)
# )
# clf_optim = optim.SGD(clf_model.parameters(), lr=1e-2, momentum=0.9)
return netG, netD, encoder, decoder, image_encoder, e_disc, clf_model
def __init__(self):
super(VAE, self).__init__()
self.fc1 = nn.Linear(imsize, 400)
self.fc21 = nn.Linear(400, 50)
self.fc22 = nn.Linear(400, 50)
self.fc3 = nn.Linear(50, 400)
self.fc4 = nn.Linear(400, imsize)
self.relu = nn.ReLU()
self.sigmoid = nn.Sigmoid()
self.netG = STAGE1_G()
# state_dict = \
# torch.load("../output/birds_stageI_2018_03_19_15_55_52/Model/netG_epoch_120.pth",
# map_location=lambda storage, loc: storage)
# self.netG.load_state_dict(state_dict)
# print('Load from: ', cfg.NET_G)
#self.image_encoder = CNN_ENCODER(128)
self.image_encoder = STAGE1_ImageEncoder()
# state_dict = \
# torch.load("../output/birds_stageI_2018_03_19_15_55_52/Model/image_encoder_epoch_120.pth",
# map_location=lambda storage, loc: storage)
# self.image_encoder.load_state_dict(state_dict)
ndf, nef = 60, 128
self.nef = nef
self.decode_lin = nn.Sequential(nn.Linear(zsize, nef * 4 * 4),
nn.BatchNorm1d(nef * 4 * 4),
nn.ReLU(True))
self.decode_img = nn.Sequential(
nn.ConvTranspose2d(nef, nef // 2, 4, 2, 1, bias=False),
nn.BatchNorm2d(nef // 2), nn.LeakyReLU(0.2, inplace=True),
nn.ConvTranspose2d(nef // 2, nef // 4, 4, 2, 1, bias=False),
nn.BatchNorm2d(nef // 4), nn.LeakyReLU(0.2, inplace=True),
nn.ConvTranspose2d(nef // 4, nef // 8, 4, 2, 1, bias=False),
nn.BatchNorm2d(nef // 8), nn.LeakyReLU(0.2, inplace=True),
nn.ConvTranspose2d(nef // 8, nef // 16, 4, 2, 1, bias=False),
nn.BatchNorm2d(nef // 16), nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(nef // 16, 3, 3, 1, 1), nn.Sigmoid())
self.encode_img = nn.Sequential(
nn.Conv2d(3, ndf, 4, 2, 1, bias=False),
nn.LeakyReLU(0.2, inplace=True),
# state size. (ndf) x 32 x 32
nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
nn.BatchNorm2d(ndf * 2),
nn.LeakyReLU(0.2, inplace=True),
# state size (ndf*2) x 16 x 16
nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
nn.BatchNorm2d(ndf * 4),
nn.LeakyReLU(0.2, inplace=True),
# state size (ndf*4) x 8 x 8
nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
nn.BatchNorm2d(ndf * 8),
# state size (ndf * 8) x 4 x 4)
nn.LeakyReLU(0.2, inplace=True)
#nn.MaxPool2d(2, stride=2),
#nn.Linear(1024, 300)
)
self.l1 = nn.Linear(480 * 4 * 4, zsize)
self.l2 = nn.Linear(480 * 4 * 4, zsize)
self.l = nn.Linear(128, zsize * 2)
self.relu = nn.ReLU()