def make_model(args, dataset):
N = args.gpu_num
class Model(): pass
model = Model()
netG = Generator(args)
netD = Discriminator(args)
netRNN = DS_SJE(args, dataset.embed_mat)
if N == 0:
setattr(model, "netG_0", netG)
setattr(model, "netD_0", netD)
setattr(model, "netRNN_0", netRNN)
else:
for i in range(N):
temp_netG = netG.copy()
temp_netD = netD.copy()
temp_netRNN = netRNN.copy()
temp_netG.gpu_id = i
temp_netD.gpu_id = i
temp_netRNN.gpu_id = i
setattr(model, f"netG_{i}", temp_netG.to_gpu(i))
setattr(model, f"netD_{i}", temp_netD.to_gpu(i))
setattr(model, f"netRNN_{i}", temp_netRNN.to_gpu(i))
return model
def build_models(conf, n_voc, seq_len):
model = dict()
netG = Generator(n_voc, conf.emb_dim, conf.hid_dim, seq_len, conf.gpu_num)
netD = Discriminator(n_voc, conf.emb_dim, conf.hid_dim, seq_len,
conf.gpu_num)
if conf.gpu_num == 0:
model["netG_0"] = netG
model["netD_0"] = netD
else:
for i in range(conf.gpu_num):
copy_netG = netG.copy()
copy_netD = netD.copy()
copy_netG.gpu_id = i
copy_netD.gpu_id = i
model[f"netG_{i}"] = copy_netG.to_gpu(i)
model[f"netD_{i}"] = copy_netD.to_gpu(i)
return model
