def train(self, dataset, attack=None):
model = self.model
folder_name = osp.join(getGitPath(), 'models')
if attack is None:
folder_name = osp.join(folder_name, 'basic_models')
targeted, attack_epochs = None, None
else:
folder_name = osp.join(folder_name, 'adversarial_models')
targeted, attack_epochs = attack.targeted, attack.attack_epochs
file_name = fileNamer(node_model=self.node_model,
dataset_name=dataset.name,
model_name=model.name,
num_layers=model.num_layers,
patience=self.patience,
seed=self.seed,
targeted=targeted,
attack_epochs=attack_epochs,
end='.pt')
model_path = osp.join(folder_name, file_name)
# load model and optimizer
if not osp.exists(model_path):
# train model
model, model_log, test_acc = self.useTrainer(data=dataset.data,
attack=attack)
torch.save((model.state_dict(), model_log, test_acc), model_path)
else:
model_state_dict, model_log, test_acc = torch.load(model_path)
model.load_state_dict(model_state_dict)
print(model_log + '\n')
self.basic_log = model_log
self.clean = test_acc
def _loadDataset(self, dataset, device):
dataset_path = osp.join(getGitPath(), 'datasets')
if dataset is DataSet.PUBMED or dataset is DataSet.CORA or dataset is DataSet.CITESEER:
dataset = Planetoid(dataset_path, dataset.string())
elif dataset is DataSet.TWITTER:
twitter_glove_path = osp.join(dataset_path, 'twitter', 'glove.pkl')
if not osp.exists(twitter_glove_path):
quit(
"Go to README and follow the download instructions to the TWITTER dataset"
)
else:
dataset = TwitterDataset(osp.dirname(twitter_glove_path))
with open(twitter_glove_path, 'rb') as file:
glove_matrix = pickle.load(file)
self.glove_matrix = torch.tensor(
glove_matrix, dtype=torch.float32).to(device)
# elif dataset is DataSet.PPI:
# PPI_path = osp.join(datasets_path, 'PPI')
# data_split = []
# for split in ('train', 'val', 'test'):
# PPI.url = 'https://data.dgl.ai/dataset/ppi.zip'
# data_split.append(PPI(PPI_path, split=split))
# elif dataset is DataSet.FLICKR:
# dataset = Flickr(osp.join(datasets_path, 'Flickr'))
data = dataset[0].to(self.device)
setattr(data, 'num_classes', dataset.num_classes)
self.num_features = data.num_features
self.num_classes = dataset.num_classes
return data
def _generateMasks(self, data, name, train_percent, val_percent):
train_mask = torch.zeros(data.num_nodes).type(torch.bool)
val_mask = torch.zeros(data.num_nodes).type(torch.bool)
test_mask = torch.zeros(data.num_nodes).type(torch.bool)
# taken from Planetoid
for c in range(data.num_classes):
idx = (data.y == c).nonzero(as_tuple=False).view(-1)
num_train_per_class = round(idx.size(0) * train_percent)
num_val_per_class = round(idx.size(0) * val_percent)
idx_permuted = idx[torch.randperm(idx.size(0))]
train_idx = idx_permuted[:num_train_per_class]
val_idx = idx_permuted[num_train_per_class:num_train_per_class +
num_val_per_class]
test_idx = idx_permuted[num_train_per_class + num_val_per_class:]
train_mask[train_idx] = True
val_mask[val_idx] = True
test_mask[test_idx] = True
print(
f'Class {c}: training: {train_idx.size(0)}, val: {val_idx.size(0)}, test: {test_idx.size(0)}'
)
masks = Masks(train_mask, val_mask, test_mask)
torch.save(masks, osp.join(getGitPath(), 'masks', name + '.dat'))
return masks
def _loadDataset(self, dataset: DataSet,
device: torch.device) -> torch_geometric.data.Data:
"""
a loader function for the requested dataset
"""
dataset_path = osp.join(getGitPath(), 'datasets')
if dataset is DataSet.PUBMED or dataset is DataSet.CORA or dataset is DataSet.CITESEER:
dataset = Planetoid(dataset_path, dataset.string())
elif dataset is DataSet.TWITTER:
twitter_glove_path = osp.join(dataset_path, 'twitter', 'glove.pkl')
if not osp.exists(twitter_glove_path):
exit(
"Go to README and follow the download instructions to the TWITTER dataset"
)
else:
dataset = TwitterDataset(osp.dirname(twitter_glove_path))
with open(twitter_glove_path, 'rb') as file:
glove_matrix = pickle.load(file)
self.glove_matrix = torch.tensor(
glove_matrix, dtype=torch.float32).to(device)
data = dataset[0].to(self.device)
setattr(data, 'num_classes', dataset.num_classes)
self.num_features = data.num_features
self.num_classes = dataset.num_classes
return data
import torch.nn.functional as F
from torch.utils.cpp_extension import load
import torch_scatter
import torch_sparse
# start of changes XXXXX
from helpers.getGitPath import getGitPath
# end of changes XXXXX
try:
try:
import kernels as custom_cuda_kernels
if not hasattr(custom_cuda_kernels, 'topk'):
raise ImportError()
except ImportError:
rgnn_dir = os.path.join(getGitPath(), 'implementation',
'model_functions', 'rgnn')
cache_dir = os.path.join(rgnn_dir, 'extension', socket.gethostname(),
torch.__version__)
csrc_dir = os.path.join(rgnn_dir, 'kernels', 'csrc')
custom_file = os.path.join(csrc_dir, 'custom.cpp')
custom_kernel_file = os.path.join(csrc_dir, 'custom_kernel.cu')
os.makedirs(cache_dir, exist_ok=True)
# start of changes XXXXX
custom_cuda_kernels = load(
name="custom_cuda",
sources=[custom_file, custom_kernel_file],
extra_cuda_cflags=['-lcusparse', '-l', 'cusparse'],
build_directory=cache_dir)
# end of changes XXXXX