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