def from_qm9_pretrained(root: str, dataset: Dataset, target: int):
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import tensorflow as tf
assert target >= 0 and target <= 12 and not target == 4
root = osp.expanduser(osp.normpath(root))
path = osp.join(root, 'pretrained_dimenet', qm9_target_dict[target])
makedirs(path)
url = f'{DimeNet.url}/{qm9_target_dict[target]}'
if not osp.exists(osp.join(path, 'checkpoint')):
download_url(f'{url}/checkpoint', path)
download_url(f'{url}/ckpt.data-00000-of-00002', path)
download_url(f'{url}/ckpt.data-00001-of-00002', path)
download_url(f'{url}/ckpt.index', path)
path = osp.join(path, 'ckpt')
reader = tf.train.load_checkpoint(path)
model = DimeNet(hidden_channels=128, out_channels=1, num_blocks=6,
num_bilinear=8, num_spherical=7, num_radial=6,
cutoff=5.0, envelope_exponent=5, num_before_skip=1,
num_after_skip=2, num_output_layers=3)
def copy_(src, name, transpose=False):
init = reader.get_tensor(f'{name}/.ATTRIBUTES/VARIABLE_VALUE')
init = torch.from_numpy(init)
if name[-6:] == 'kernel':
init = init.t()
src.data.copy_(init)
copy_(model.rbf.freq, 'rbf_layer/frequencies')
copy_(model.emb.emb.weight, 'emb_block/embeddings')
copy_(model.emb.lin_rbf.weight, 'emb_block/dense_rbf/kernel')
copy_(model.emb.lin_rbf.bias, 'emb_block/dense_rbf/bias')
copy_(model.emb.lin.weight, 'emb_block/dense/kernel')
copy_(model.emb.lin.bias, 'emb_block/dense/bias')
for i, block in enumerate(model.output_blocks):
copy_(block.lin_rbf.weight, f'output_blocks/{i}/dense_rbf/kernel')
for j, lin in enumerate(block.lins):
copy_(lin.weight, f'output_blocks/{i}/dense_layers/{j}/kernel')
copy_(lin.bias, f'output_blocks/{i}/dense_layers/{j}/bias')
copy_(block.lin.weight, f'output_blocks/{i}/dense_final/kernel')
for i, block in enumerate(model.interaction_blocks):
copy_(block.lin_rbf.weight, f'int_blocks/{i}/dense_rbf/kernel')
copy_(block.lin_sbf.weight, f'int_blocks/{i}/dense_sbf/kernel')
copy_(block.lin_kj.weight, f'int_blocks/{i}/dense_kj/kernel')
copy_(block.lin_kj.bias, f'int_blocks/{i}/dense_kj/bias')
copy_(block.lin_ji.weight, f'int_blocks/{i}/dense_ji/kernel')
copy_(block.lin_ji.bias, f'int_blocks/{i}/dense_ji/bias')
copy_(block.W, f'int_blocks/{i}/bilinear')
for j, layer in enumerate(block.layers_before_skip):
copy_(layer.lin1.weight,
f'int_blocks/{i}/layers_before_skip/{j}/dense_1/kernel')
copy_(layer.lin1.bias,
f'int_blocks/{i}/layers_before_skip/{j}/dense_1/bias')
copy_(layer.lin2.weight,
f'int_blocks/{i}/layers_before_skip/{j}/dense_2/kernel')
copy_(layer.lin2.bias,
f'int_blocks/{i}/layers_before_skip/{j}/dense_2/bias')
copy_(block.lin.weight, f'int_blocks/{i}/final_before_skip/kernel')
copy_(block.lin.bias, f'int_blocks/{i}/final_before_skip/bias')
for j, layer in enumerate(block.layers_after_skip):
copy_(layer.lin1.weight,
f'int_blocks/{i}/layers_after_skip/{j}/dense_1/kernel')
copy_(layer.lin1.bias,
f'int_blocks/{i}/layers_after_skip/{j}/dense_1/bias')
copy_(layer.lin2.weight,
f'int_blocks/{i}/layers_after_skip/{j}/dense_2/kernel')
copy_(layer.lin2.bias,
f'int_blocks/{i}/layers_after_skip/{j}/dense_2/bias')
# Use the same random seed as the official DimeNet` implementation.
random_state = np.random.RandomState(seed=42)
perm = torch.from_numpy(random_state.permutation(np.arange(130831)))
train_idx = perm[:110000]
val_idx = perm[110000:120000]
test_idx = perm[120000:]
return model, (dataset[train_idx], dataset[val_idx], dataset[test_idx])
def download(self):
path = download_url('{}/{}.zip'.format(self.url, self.name), self.root)
extract_zip(path, self.root)
os.unlink(path)
os.rename(osp.join(self.root, self.name), self.raw_dir)
def download(self):
url = f'{self.url}/{self.file_names[self.name]}'
path = download_url(url, self.raw_dir)
if url[-4:] == '.zip':
extract_zip(path, self.raw_dir)
os.unlink(path)
def download(self):
for file_name in self.raw_file_names:
download_url(MD17.raw_url + file_name, self.raw_dir)
def download(self):
download_url(self.url, self.raw_dir)
def download(self):
for url in self.raw_url:
download_url(url, self.raw_dir)
def download(self):
for filename in self.raw_file_names:
download_url(f'{self.url}/{filename}', self.raw_dir)
def download(self):
for name in self.raw_file_names:
download_url('{}/{}'.format(self.url, name), self.raw_dir)
def download(self):
download_url(osp.join(self.url, self.name + '.npz'), self.raw_dir)
def download(self):
path = download_url(self.url.format(self.file_id), self.root)
extract_zip(path, self.root)
os.unlink(path)
shutil.rmtree(self.raw_dir)
os.rename(osp.join(self.root, 'DBP15K'), self.raw_dir)
def download(self):
path = download_url(self.url, self.root)
extract_zip(path, self.root)
os.unlink(path)
name = self.__class__.__name__.lower()
os.rename(osp.join(self.root, name), self.raw_dir)
def download(self):
path = download_url(self.url, self.root)
extract_zip(path, self.root)
shutil.rmtree(self.raw_dir)
os.rename(osp.join(self.root, 'PF-dataset-PASCAL'), self.raw_dir)
def download(self):
for part in ['edges', 'features', 'target']:
download_url(f'{self.url}/{self.name}/{part}.csv', self.raw_dir)
def download(self):
download_url(f'{self.url}/{self.name}.npz', self.raw_dir)
def download(self):
path = download_url(self.urls[self.name], self.root)
extract_zip(path, self.root)
os.unlink(path)
folder = osp.join(self.root, 'ModelNet{}'.format(self.name))
os.rename(folder, self.raw_dir)
def download(self):
download_url(self.url + self.raw_file_names, self.raw_dir)
def download(self):
for url in self.urls:
path = download_url(url, self.raw_dir)
extract_gz(path, self.raw_dir)
os.unlink(path)
def download(self):
path = download_url(self.url, self.raw_dir)
extract_tar(path, self.raw_dir, mode='r')
os.unlink(path)
def test_to_superpixels():
root = osp.join('/', 'tmp', str(random.randrange(sys.maxsize)))
raw_folder = osp.join(root, 'MNIST', 'raw')
processed_folder = osp.join(root, 'MNIST', 'processed')
makedirs(raw_folder)
makedirs(processed_folder)
for resource in resources:
path = download_url(resource, raw_folder)
extract_gz(path, osp.join(root, raw_folder))
test_set = (
read_image_file(osp.join(raw_folder, 't10k-images-idx3-ubyte')),
read_label_file(osp.join(raw_folder, 't10k-labels-idx1-ubyte')),
)
torch.save(test_set, osp.join(processed_folder, 'training.pt'))
torch.save(test_set, osp.join(processed_folder, 'test.pt'))
dataset = MNIST(root, download=False)
dataset.transform = T.Compose([T.ToTensor(), ToSLIC()])
data, y = dataset[0]
assert len(data) == 2
assert data.pos.dim() == 2 and data.pos.size(1) == 2
assert data.x.dim() == 2 and data.x.size(1) == 1
assert data.pos.size(0) == data.x.size(0)
assert y == 7
loader = DataLoader(dataset, batch_size=2, shuffle=False)
for data, y in loader:
assert len(data) == 4
assert data.pos.dim() == 2 and data.pos.size(1) == 2
assert data.x.dim() == 2 and data.x.size(1) == 1
assert data.batch.dim() == 1
assert data.ptr.dim() == 1
assert data.pos.size(0) == data.x.size(0) == data.batch.size(0)
assert y.tolist() == [7, 2]
break
dataset.transform = T.Compose(
[T.ToTensor(), ToSLIC(add_seg=True, add_img=True)])
data, y = dataset[0]
assert len(data) == 4
assert data.pos.dim() == 2 and data.pos.size(1) == 2
assert data.x.dim() == 2 and data.x.size(1) == 1
assert data.pos.size(0) == data.x.size(0)
assert data.seg.size() == (1, 28, 28)
assert data.img.size() == (1, 1, 28, 28)
assert data.seg.max().item() + 1 == data.x.size(0)
assert y == 7
loader = DataLoader(dataset, batch_size=2, shuffle=False)
for data, y in loader:
assert len(data) == 6
assert data.pos.dim() == 2 and data.pos.size(1) == 2
assert data.x.dim() == 2 and data.x.size(1) == 1
assert data.batch.dim() == 1
assert data.ptr.dim() == 1
assert data.pos.size(0) == data.x.size(0) == data.batch.size(0)
assert data.seg.size() == (2, 28, 28)
assert data.img.size() == (2, 1, 28, 28)
assert y.tolist() == [7, 2]
break
shutil.rmtree(root)
def download(self):
download_url(self.url, self.raw_dir)
os.system('gunzip -f ' +
os.path.join(self.raw_dir, self.raw_file_names + '.gz'))
def download(self):
url = self.url.format(self.names[self.name][1])
path = download_url(url, self.raw_dir)
if self.names[self.name][1][-2:] == 'gz':
extract_gz(path, self.raw_dir)
os.unlink(path)
def download(self):
download_url(self.url.format(self.name), self.raw_dir)
def download(self):
path = download_url(self.url, self.raw_dir)
extract_zip(path, self.raw_dir)
os.unlink(path)
def download(self):
url = self.csl_url
url = f'{self.url}/{self.name}.zip' if self.name != 'CSL' else url
path = download_url(url, self.raw_dir)
extract_zip(path, self.raw_dir)
os.unlink(path)
def download(self):
url = self.processed_url if rdkit is None else self.raw_url
file_path = download_url(url, self.raw_dir)
extract_zip(file_path, self.raw_dir)
os.unlink(file_path)
def download(self):
path = download_url(self.url.format(self.name), self.root)
extract_tar(path, self.raw_dir)
os.unlink(path)
def download(self):
path = download_url(self.url, self.raw_dir)
extract_zip(path, self.raw_dir)
os.remove(path)
def download(self):
for name in self.raw_file_names:
url = '{}/{}.zip'.format(self.url, name)
path = download_url(url, self.raw_dir)
extract_zip(path, self.raw_dir)
os.unlink(path)
def download(self):
context = ssl._create_default_https_context
ssl._create_default_https_context = ssl._create_unverified_context
download_url(f'{self.url}/{self.raw_file_names}', self.raw_dir)
ssl._create_default_https_context = context
def from_qm9_pretrained(root: str, dataset: Dataset, target: int):
import ase
import schnetpack as spk # noqa
assert target >= 0 and target <= 12
units = [1] * 12
units[0] = ase.units.Debye
units[1] = ase.units.Bohr**3
units[5] = ase.units.Bohr**2
root = osp.expanduser(osp.normpath(root))
makedirs(root)
folder = 'trained_schnet_models'
if not osp.exists(osp.join(root, folder)):
path = download_url(SchNet.url, root)
extract_zip(path, root)
os.unlink(path)
name = f'qm9_{qm9_target_dict[target]}'
path = osp.join(root, 'trained_schnet_models', name, 'split.npz')
split = np.load(path)
train_idx = split['train_idx']
val_idx = split['val_idx']
test_idx = split['test_idx']
# Filter the splits to only contain characterized molecules.
idx = dataset.data.idx
assoc = idx.new_empty(idx.max().item() + 1)
assoc[idx] = torch.arange(idx.size(0))
train_idx = assoc[train_idx[np.isin(train_idx, idx)]]
val_idx = assoc[val_idx[np.isin(val_idx, idx)]]
test_idx = assoc[test_idx[np.isin(test_idx, idx)]]
path = osp.join(root, 'trained_schnet_models', name, 'best_model')
with warnings.catch_warnings():
warnings.simplefilter('ignore')
state = torch.load(path, map_location='cpu')
net = SchNet(hidden_channels=128,
num_filters=128,
num_interactions=6,
num_gaussians=50,
cutoff=10.0,
atomref=dataset.atomref(target))
net.embedding.weight = state.representation.embedding.weight
for int1, int2 in zip(state.representation.interactions,
net.interactions):
int2.mlp[0].weight = int1.filter_network[0].weight
int2.mlp[0].bias = int1.filter_network[0].bias
int2.mlp[2].weight = int1.filter_network[1].weight
int2.mlp[2].bias = int1.filter_network[1].bias
int2.lin.weight = int1.dense.weight
int2.lin.bias = int1.dense.bias
int2.conv.lin1.weight = int1.cfconv.in2f.weight
int2.conv.lin2.weight = int1.cfconv.f2out.weight
int2.conv.lin2.bias = int1.cfconv.f2out.bias
net.lin1.weight = state.output_modules[0].out_net[1].out_net[0].weight
net.lin1.bias = state.output_modules[0].out_net[1].out_net[0].bias
net.lin2.weight = state.output_modules[0].out_net[1].out_net[1].weight
net.lin2.bias = state.output_modules[0].out_net[1].out_net[1].bias
mean = state.output_modules[0].atom_pool.average
net.readout = 'mean' if mean is True else 'add'
dipole = state.output_modules[0].__class__.__name__ == 'DipoleMoment'
net.dipole = dipole
net.mean = state.output_modules[0].standardize.mean.item()
net.std = state.output_modules[0].standardize.stddev.item()
if state.output_modules[0].atomref is not None:
net.atomref.weight = state.output_modules[0].atomref.weight
else:
net.atomref = None
net.scale = 1. / units[target]
return net, (dataset[train_idx], dataset[val_idx], dataset[test_idx])
