# write the origin model again.
# D0 = model.E_Net_small()
encoder_model = torch.load("/home/bike/2027/generative-models/tools/D_mnist_little3.model")
# 24*24
enet = model.E_Net_small().cuda()
enet.load_state_dict(encoder_model)
add_in_feature = feature_size+ hidden_d # Add one dimension data for the input_feature data.
G = model.G_Net_FM_3(ngpu,add_in_feature,main_gpu=main_gpu).cuda()
# g_model = torch.load("./fm21/G_95000.model")
# G.load_state_dict(g_model)
in_channel = 1
E = model.Ev_Net_conv(ngpu,in_channel,main_gpu=main_gpu).cuda()
# E.apply(weights_init)
G.apply(weights_init)
d_in_demension = 1
D = model.D_Net_conv(ngpu,d_in_demension,main_gpu=main_gpu).cuda()
D.apply(weights_init)
G_solver = optim.Adam(G.parameters(),lr = 1e-4)
D_solver = optim.Adam(D.parameters(), lr = 1e-4)
E_solver = optim.Adam(E.parameters(), lr=1e-4)
half_label = Variable(torch.ones(batch_size)*0.5).cuda()
check_points = 500
num_epoches = 100000
sys.stdout = mutil.Logger(out_dir)
# else:
# print("you have already creat one.")
# exit(1)
#
#
# def xavier_init(size):
# in_dim = size[0]
# xavier_stddev = 1. / np.sqrt(in_dim / 2.)
# return Variable(torch.randn(*size) * xavier_stddev, requires_grad=True)
in_channel = 2
G = model.G_Net_conv(ngpu).cuda()
D = model.D_Net_conv(ngpu, 1).cuda()
E = model.Ev_Net_conv(ngpu, 1).cuda()
"""Weight Initialization"""
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
m.weight.data.normal_(0.0, 0.02)
# elif classname.find('BatchNorm') != -1:
# m.weight.data.normal_(1.0, 0.02)
# m.bias.data.fill_(0)
G.apply(weights_init)
D.apply(weights_init)
E.apply(weights_init)
