def __init__(self, device, cfg):
self.device = device
self.cfg = cfg
self.bs = self.cfg.getint('training', 'batchsize')
#Data pipeline
self.data = Data(cfg, save=False)
ds_train = DS(self.data, cfg)
if len(ds_train) != 0:
self.loader_train = DataLoader(
ds_train,
batch_size=self.bs,
shuffle=True,
drop_last=False,
pin_memory=True,
num_workers=0,
)
else:
self.loader_train = []
self.steps_per_epoch = int(
len(self.loader_train) /
(self.cfg.getint('training', 'n_critic') + 1))
print(len(self.loader_train), self.steps_per_epoch)
self.ff_inp = self.data.ff_inp
self.ff_out = self.data.ff_out
if len(ds_train) != 0:
ds_val = DS(self.data, cfg, train=False)
if len(ds_val) != 0:
self.loader_val = DataLoader(
ds_val,
batch_size=self.bs,
shuffle=True,
drop_last=False,
pin_memory=True,
num_workers=0,
)
else:
self.loader_val = []
#model
self.name = cfg.get('model', 'name')
self.cond = cfg.getboolean('model', 'cond')
if self.cond and not self.data.pairs:
raise Exception(
'conditional GAN can only be used with pairs of snapshots for both resolutions.'
)
self.out_env = cfg.getboolean('model', 'out_env')
self.n_env_mols = cfg.getint('universe', 'n_env_mols')
self.n_mols = cfg.getint('universe', 'n_mols')
self.feature_dim = self.ff_inp.n_atom_chns
self.feature_dim += self.ff_inp.n_atom_chns
"""
if self.n_env_mols != 0:
self.feature_dim += self.ff_inp.n_atom_chns
if self.out_env:
self.feature_dim += self.ff_out.n_atom_chns
"""
self.target_dim = self.ff_out.n_atoms * self.n_mols
self.critic_dim = self.ff_inp.n_atom_chns + self.target_dim
#if self.cond:
#self.critic_dim += self.feature_dim
self.z_dim = int(cfg.getint('model', 'noise_dim'))
#print(self.ff_inp.n_atom_chns, self.n_input, self.n_out, self.ff_out.n_atoms)
self.step = 0
self.epoch = 0
# Make Dirs for saving
self.out = OutputHandler(
self.name,
self.cfg.getint('training', 'n_checkpoints'),
self.cfg.get('model', 'output_dir'),
)
self.energy_inp = Energy_torch(self.ff_inp, self.device)
self.energy_out = Energy_torch(self.ff_out, self.device)
self.ol_weight = cfg.getfloat('prior', 'ol')
prior_weights = self.cfg.get('prior', 'weights')
self.prior_weights = [float(v) for v in prior_weights.split(",")]
prior_schedule = self.cfg.get('prior', 'schedule')
try:
self.prior_schedule = np.array(
[0] + [int(v) for v in prior_schedule.split(",")])
except:
self.prior_schedule = [0]
self.ratio_bonded_nonbonded = cfg.getfloat('prior',
'ratio_bonded_nonbonded')
self.prior_mode = cfg.get('prior', 'mode')
print(self.prior_mode)
#Model selection
#if cfg.get('model', 'model_type') == "tiny":
# print("Using tiny model")
if self.z_dim != 0:
self.generator = model.G_tiny_with_noise(
z_dim=self.z_dim,
n_input=self.feature_dim,
n_output=self.target_dim,
start_channels=self.cfg.getint('model', 'n_chns'),
fac=1,
sn=self.cfg.getint('model', 'sn_gen'),
device=device)
print("Using tiny generator with noise")
else:
self.generator = model.G_tiny(
n_input=self.feature_dim,
n_output=self.target_dim,
start_channels=self.cfg.getint('model', 'n_chns'),
fac=1,
sn=self.cfg.getint('model', 'sn_gen'),
device=device)
print("Using tiny generator without noise")
if cfg.getint('grid', 'resolution') == 8:
self.critic = model.C_tiny_mbd(
in_channels=self.critic_dim,
start_channels=self.cfg.getint('model', 'n_chns'),
fac=1,
sn=self.cfg.getint('model', 'sn_crit'),
device=device)
print("Using tiny critic with resolution 8")
else:
self.critic = model.C_tiny16(
in_channels=self.critic_dim,
start_channels=self.cfg.getint('model', 'n_chns'),
fac=1,
sn=self.cfg.getint('model', 'sn_crit'),
device=device)
print("Using tiny critic with resolution 16")
self.use_gp = cfg.getboolean('model', 'gp')
self.use_ol = cfg.getboolean('model', 'ol')
self.use_energy = cfg.getboolean('model', 'energy')
self.critic.to(device=device)
self.generator.to(device=device)
#self.mse.to(device=device)
lr_gen = cfg.getfloat('training', 'lr_gen')
lr_crit = cfg.getfloat('training', 'lr_crit')
#self.opt_generator_pretrain = Adam(self.generator.parameters(), lr=lr_gen, betas=(0, 0.9))
self.opt_generator = Adam(self.generator.parameters(),
lr=lr_gen,
betas=(0, 0.9))
self.opt_critic = Adam(self.critic.parameters(),
lr=lr_crit,
betas=(0, 0.9))
self.restored_model = False
self.restore_latest_checkpoint()
def __init__(self, device, cfg):
self.device = device
self.cfg = cfg
self.bs = self.cfg.getint('training', 'batchsize')
#Data pipeline
self.data = Data(cfg, save=False)
ds_train = DS(self.data, cfg)
if len(ds_train) != 0:
self.loader_train = DataLoader(
ds_train,
batch_size=self.bs,
shuffle=True,
drop_last=False,
pin_memory=True,
num_workers=0,
)
else:
self.loader_train = []
self.steps_per_epoch = int(len(self.loader_train) / (self.cfg.getint('training', 'n_critic') + 1))
print(len(self.loader_train), self.steps_per_epoch)
self.ff_inp = self.data.ff_inp
self.ff_out = self.data.ff_out
if len(ds_train) != 0:
ds_val = DS(self.data, cfg, train=False)
if len(ds_val) != 0:
self.loader_val = DataLoader(
ds_val,
batch_size=self.bs,
shuffle=True,
drop_last=False,
pin_memory=True,
num_workers=0,
)
else:
self.loader_val = []
#model
self.name = cfg.get('model', 'name')
self.cond = cfg.getboolean('model', 'cond')
if self.cond and not self.data.pairs:
raise Exception('conditional GAN can only be used with pairs of snapshots for both resolutions.')
self.out_env = cfg.getboolean('model', 'out_env')
self.n_env_mols = cfg.getint('universe', 'n_env_mols')
self.feature_dim = self.ff_inp.n_atom_chns
if self.n_env_mols != 0:
self.feature_dim += self.ff_inp.n_atom_chns
if self.out_env:
self.feature_dim += self.ff_out.n_atom_chns
self.target_dim = self.ff_out.n_atoms
self.critic_dim = self.target_dim
if self.cond:
self.critic_dim += self.feature_dim
self.z_dim = int(cfg.getint('model', 'noise_dim'))
self.recon = cfg.getboolean('model', 'recon')
if self.recon and self.ff_out.n_atoms < self.ff_inp.n_atoms:
raise Exception("reconstruction error only applicable when going from lower to higher resolution")
elif self.recon:
self.mapping = {}
self.recon_weight = cfg.getfloat('prior', 'rec')
map_file = self.data.dir_mapping / self.cfg.get('model', 'map_file')
if map_file.exists():
for line in read_between("[map]", "[/map]", self.data.dir_mapping / self.cfg.get('model', 'map_file')):
if len(line.split()) == 2:
out_ndx = int(line.split()[0]) - 1
inp_ndx = int(line.split()[1]) - 1
self.mapping[out_ndx] = inp_ndx
print(self.mapping)
if len(self.mapping) != self.ff_out.n_atoms:
raise Exception("something wrong with the mapping file")
inp_masses, out_masses = np.zeros(self.ff_inp.n_atoms), np.zeros(self.ff_out.n_atoms)
for k in range(0, self.ff_out.n_atoms):
out_masses[k] = self.data.samples_train_out[0].mols[0].atoms[k].type.mass
inp_masses[self.mapping[k]] += out_masses[k]
#for a in self.data.samples_train_out[0].mols[0].atoms:
# out_masses.append(a.type.mass)
self.inp_masses = torch.from_numpy(inp_masses).to(self.device).float()
self.out_masses = torch.from_numpy(out_masses).to(self.device).float()
self.inp_masses = self.inp_masses[None, :, None]
self.out_masses = self.out_masses[None, :, None]
else:
raise Exception("no mapping file but training with reconstruction error")
#print(self.ff_inp.n_atom_chns, self.n_input, self.n_out, self.ff_out.n_atoms)
self.step = 0
self.epoch = 0
# Make Dirs for saving
self.out = OutputHandler(
self.name,
self.cfg.getint('training', 'n_checkpoints'),
self.cfg.get('model', 'output_dir'),
)
self.energy_inp = Energy_torch(self.ff_inp, self.device)
self.energy_out = Energy_torch(self.ff_out, self.device)
self.ol_weight = cfg.getfloat('prior', 'ol')
prior_weights = self.cfg.get('prior', 'weights')
self.prior_weights = [float(v) for v in prior_weights.split(",")]
prior_schedule = self.cfg.get('prior', 'schedule')
try:
self.prior_schedule = np.array([0] + [int(v) for v in prior_schedule.split(",")])
except:
self.prior_schedule = [0]
self.ratio_bonded_nonbonded = cfg.getfloat('prior', 'ratio_bonded_nonbonded')
self.prior_mode = cfg.get('prior', 'mode')
print(self.prior_mode)
#Model selection
#if cfg.get('model', 'model_type') == "tiny":
# print("Using tiny model")
if self.z_dim != 0:
self.generator = model.G_tiny_with_noise(z_dim=self.z_dim,
n_input=self.feature_dim,
n_output=self.target_dim,
start_channels=self.cfg.getint('model', 'n_chns'),
fac=1,
sn=self.cfg.getint('model', 'sn_gen'),
device=device)
print("Using tiny generator with noise")
else:
self.generator = model.G_tiny(n_input=self.feature_dim,
n_output=self.target_dim,
start_channels=self.cfg.getint('model', 'n_chns'),
fac=1,
sn=self.cfg.getint('model', 'sn_gen'),
device=device)
print("Using tiny generator without noise")
if cfg.getint('grid', 'resolution') == 8:
self.critic = model.C_tiny_mbd(in_channels=self.critic_dim,
start_channels=self.cfg.getint('model', 'n_chns'),
fac=1, sn=self.cfg.getint('model', 'sn_crit'),
device=device)
print("Using tiny critic with resolution 8")
else:
self.critic = model.C_tiny16(in_channels=self.critic_dim,
start_channels=self.cfg.getint('model', 'n_chns'),
fac=1, sn=self.cfg.getint('model', 'sn_crit'),
device=device)
print("Using tiny critic with resolution 16")
self.use_gp = cfg.getboolean('model', 'gp')
self.use_ol = cfg.getboolean('model', 'ol')
self.use_energy = cfg.getboolean('model', 'energy')
self.critic.to(device=device)
self.generator.to(device=device)
#self.mse.to(device=device)
lr_gen = cfg.getfloat('training', 'lr_gen')
lr_crit = cfg.getfloat('training', 'lr_crit')
#self.opt_generator_pretrain = Adam(self.generator.parameters(), lr=lr_gen, betas=(0, 0.9))
self.opt_generator = Adam(self.generator.parameters(), lr=lr_gen, betas=(0, 0.9))
self.opt_critic = Adam(self.critic.parameters(), lr=lr_crit, betas=(0, 0.9))
self.restored_model = False
self.restore_latest_checkpoint()