def get_train_val_loader(path,
train_batch_size=64,
val_batch_size=64,
val_split=1 / 12):
train_dataset = MNISTSuperpixels(path, "train", transform=T.Cartesian())
dataset_size = len(train_dataset)
indices = list(range(dataset_size))
split = int(val_split * dataset_size)
np.random.seed(43)
np.random.shuffle(indices)
train_indices = indices[split:]
val_indices = indices[:split]
train_sampler = SubsetRandomSampler(train_indices)
val_sampler = SubsetRandomSampler(val_indices)
validation_loader = DataLoader(train_dataset,
batch_size=val_batch_size,
sampler=val_sampler,
shuffle=False)
train_loader = DataLoader(train_dataset,
batch_size=train_batch_size,
sampler=train_sampler,
shuffle=False)
return train_loader, validation_loader
def get_data():
dataset = args.name
path = '../data/geometric/MNIST'
trainset = MNISTSuperpixels(path, train=True)
testset = MNISTSuperpixels(path, train=False)
lenTrain = len(trainset)
lenTest = len(testset)
trainLoader = DataLoader(trainset[:lenTrain // 50],
batch_size=1,
shuffle=False)
testloader = DataLoader(testset[:lenTest // 50],
batch_size=1,
shuffle=False)
return trainLoader, testloader
def load_pyg(name, dataset_dir):
"""
Load PyG dataset objects. (More PyG datasets will be supported)
Args:
name (string): dataset name
dataset_dir (string): data directory
Returns: PyG dataset object
"""
dataset_dir = '{}/{}'.format(dataset_dir, name)
if name in ['Cora', 'CiteSeer', 'PubMed']:
dataset = Planetoid(dataset_dir, name)
elif name[:3] == 'TU_':
# TU_IMDB doesn't have node features
if name[3:] == 'IMDB':
name = 'IMDB-MULTI'
dataset = TUDataset(dataset_dir, name, transform=T.Constant())
else:
dataset = TUDataset(dataset_dir, name[3:])
elif name == 'Karate':
dataset = KarateClub()
elif 'Coauthor' in name:
if 'CS' in name:
dataset = Coauthor(dataset_dir, name='CS')
else:
dataset = Coauthor(dataset_dir, name='Physics')
elif 'Amazon' in name:
if 'Computers' in name:
dataset = Amazon(dataset_dir, name='Computers')
else:
dataset = Amazon(dataset_dir, name='Photo')
elif name == 'MNIST':
dataset = MNISTSuperpixels(dataset_dir)
elif name == 'PPI':
dataset = PPI(dataset_dir)
elif name == 'QM7b':
dataset = QM7b(dataset_dir)
else:
raise ValueError('{} not support'.format(name))
return dataset
def GCN(dataset, params, num_pre_epochs, num_epochs, MonteSize, PruningTimes,
width, lr, savepath):
Batch_size = int(params[0])
for Monte_iter in range(MonteSize):
# Data
best_loss = float('inf') # best test loss
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
TrainConvergence = []
TestConvergence = []
# model
root = '/git/data/GraphData/' + dataset
if dataset == 'Cora':
model_name = "PruningGCN"
datasetroot = Planetoid(root=root, name=dataset).shuffle()
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
testloader = DataListLoader(datasetroot,
batch_size=100,
shuffle=False)
model_to_save = './checkpoint/{}-{}-param_{}_{}_{}_{}-Monte_{}-ckpt.pth'.format(
dataset, model_name, params[0], params[1], params[2],
params[3], Monte_iter)
if Monte_iter == 0:
if resume == True and os.path.exists(model_to_save):
[
OptimizedNet, NewNetworksize, TrainConvergence,
TestConvergence, start_epoch
] = ResumeModel(model_to_save)
if start_epoch >= num_epochs - 1:
continue
else:
net = Net(datasetroot, width)
#net.apply(weight_reset)
elif dataset == 'ENZYMES' or dataset == 'MUTAG':
model_name = "topk_pool_Net"
datasetroot = TUDataset(root, name=dataset)
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
testloader = DataListLoader(datasetroot,
batch_size=100,
shuffle=False)
model_to_save = './checkpoint/{}-{}-param_{}_{}_{}_{}-ckpt.pth'.format(
dataset, model_name, params[0], params[1], params[2],
params[3])
if Monte_iter == 0:
if resume == True and os.path.exists(model_to_save):
[
OptimizedNet, NewNetworksize, TrainConvergence,
TestConvergence, start_epoch
] = ResumeModel(model_to_save)
if start_epoch >= num_epochs - 1:
continue
else:
net = topk_pool_Net(datasetroot, width)
elif dataset == 'MNIST':
datasetroot = MNISTSuperpixels(root=root,
transform=T.Cartesian()).shuffle()
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
testloader = DataListLoader(datasetroot,
batch_size=100,
shuffle=False)
model_name = 'SPlineNet'
model_to_save = './checkpoint/{}-{}-param_{}_{}_{}_{}-Mon_{}-ckpt.pth'.format(
dataset, model_name, params[0], params[1], params[2],
params[3], Monte_iter)
if resume == "True" and os.path.exists(model_to_save):
[net, TrainConvergence, TestConvergence,
start_epoch] = ResumeModel(model_to_save)
if start_epoch >= num_epochs - 1:
continue
else:
#net=Net(datasetroot,width)
net = SPlineNet(datasetroot, width)
elif dataset == 'CIFAR10':
if resume == "True" and os.path.exists(model_to_save):
[net, TrainConvergence, TestConvergence,
start_epoch] = ResumeModel(model_to_save)
if start_epoch >= num_epochs - 1:
continue
else:
net = getattr(CIFAR10_resnet, 'Resnet20_CIFAR10')(params[1])
else:
raise Exception(
"The dataset is:{}, it isn't existed.".format(dataset))
if Monte_iter == 0 and start_epoch == 0:
print('Let\'s use', torch.cuda.device_count(), 'GPUs!')
net = DataParallel(net)
device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
net = net.to(device)
#cudnn.benchmark = True
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=5e-4)
logging(
'Batch size: {}, Number of layers:{} ConCoeff: {}, CutoffCoffi:{}, MonteSize:{}'
.format(params[0], params[1], params[2], params[3],
Monte_iter))
for epoch in range(num_pre_epochs):
PreTrainLoss = train(trainloader, net, optimizer, criterion)
print('\nEpoch: {}, Average pre-tain loss: {:.4f} \n'.format(
epoch, PreTrainLoss[0]))
NewNetworksize = RetainNetworkSize(net, params[2])
del net
#NewNetworksize=width
for pruningIter in range(PruningTimes):
if pruningIter > 0:
[
OptimizedNet, NewNetworksize, TrainConvergence,
TestConvergence, start_epoch
] = ResumeModel(model_to_save)
elif dataset == 'Cora' and start_epoch == 0:
OptimizedNet = Net(datasetroot, NewNetworksize[0:-1])
elif dataset == 'ENZYMES' and start_epoch == 0:
NewNetworkSizeAdjust = NewNetworksize[0:-1]
NewNetworkSizeAdjust[0] = width[0] - 1
OptimizedNet = topk_pool_Net(datasetroot, NewNetworkSizeAdjust)
OptimizedNet = DataParallel(OptimizedNet)
device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
OptimizedNet = OptimizedNet.to(device)
cudnn.benchmark = True
criterionNew = nn.CrossEntropyLoss()
optimizerNew = optim.SGD(OptimizedNet.parameters(),
lr=lr,
momentum=0.9,
weight_decay=5e-4)
criterion = nn.CrossEntropyLoss()
for epoch in range(start_epoch, num_epochs):
TrainLoss = train(trainloader, OptimizedNet, optimizerNew,
criterionNew)
print('\n Epoch: {}, Average tain loss: {:.4f} \n'.format(
epoch, TrainLoss[0]))
TrainConvergence.append(statistics.mean(TrainLoss))
NewNetworksize = RetainNetworkSize(OptimizedNet, params[2])
TestConvergence.append(
statistics.mean(test(testloader, OptimizedNet, criterion)))
# save model
if TestConvergence[epoch] < best_loss:
logging('Saving..')
state = {
'net': OptimizedNet.module,
'TrainConvergence': TrainConvergence,
'TestConvergence': TestConvergence,
'epoch': num_epochs,
'NewNetworksize': NewNetworksize[0:-1],
}
if not os.path.isdir('checkpoint'):
os.mkdir('checkpoint')
torch.save(state, model_to_save)
best_loss = TestConvergence[epoch]
## save recurrence plots
"""if epoch%20==0:
save_recurrencePlots_file="../Results/RecurrencePlots/RecurrencePlots_{}_{}_BatchSize{} \_ConCoeffi{}_epoch{}.png".format(dataset, model_name,params[0],params[1],epoch)
save_recurrencePlots(net,save_recurrencePlots_file)"""
FileName = "{}-{}-param_{}_{}_{}_{}-monte_{}".format(
dataset, model_name, params[0], params[1], params[2], params[3],
Monte_iter)
np.save(savepath + 'TrainConvergence-' + FileName, TrainConvergence)
#np.save(savepath+'TestConvergence-'+FileName,TestConvergence)
#torch.cuda.empty_cache()
print_nvidia_useage()
if return_output == True:
return TestConvergence[-1], net.module.fc.weight
else:
pass
import os.path as osp
import torch
import torch.nn.functional as F
from torch_geometric.datasets import MNISTSuperpixels
import torch_geometric.transforms as T
from torch_geometric.data import DataLoader
from torch_geometric.nn import SplineConv, voxel_grid, max_pool, max_pool_x
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'MNIST')
transform = T.Cartesian(cat=False)
train_dataset = MNISTSuperpixels(path, True, transform=transform)
test_dataset = MNISTSuperpixels(path, False, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64)
d = train_dataset
class Net(torch.nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = SplineConv(d.num_features, 32, dim=2, kernel_size=5)
self.conv2 = SplineConv(32, 64, dim=2, kernel_size=5)
self.conv3 = SplineConv(64, 64, dim=2, kernel_size=5)
self.fc1 = torch.nn.Linear(4 * 64, 128)
self.fc2 = torch.nn.Linear(128, d.num_classes)
def forward(self, data):
data.x = F.elu(self.conv1(data.x, data.edge_index, data.edge_attr))
cluster = voxel_grid(data.pos, data.batch, size=5, start=0, end=28)
data.edge_attr = None
def prepare_data(self):
path = osp.join(
osp.dirname(osp.realpath(__file__)), "..", "..", "data", self.NAME
)
self.train_dataset = MNISTSuperpixels(path, True, transform=self._transform)
self.test_dataset = MNISTSuperpixels(path, False, transform=self._transform)
# rr, cc = draw.line(start_y, start_x, end_y, end_x)
# image[rr, cc] = [1, 0, 0]
# # Draw nodes
# for i in range(position.size(0)):
# x, y = position[i]
# rr, cc = draw.circle(y, x, 10, shape=shape)
# image[rr, cc] = [1, 0, 0]
# rr, cc = draw.circle(y, x, 9, shape=shape)
# image[rr, cc] = [input[i], input[i], input[i]]
# return image
image_dataset = MNIST('/tmp/MNIST', train=True, download=True)
graph_dataset = MNISTSuperpixels('~/MNISTSuperpixels', train=True)
scale = 32
rescale = 4
offset = torch.FloatTensor([1, 1])
example = 35
image, _ = image_dataset[example]
image = np.array(image)
data = graph_dataset[example]
input, adj, position = data['input'], data['adj'], data['position']
print(data['target'])
# adj = grid(torch.Size([28, 28]), connectivity=8)
# position = grid_position(torch.Size([28, 28]))
# input = image.flatten() / 255.0
def get_dataset(name, sparse=True, feat_str="deg+ak3+reall", root=None):
if root is None or root == '':
path = osp.join(osp.expanduser('~'), 'pyG_data', name)
else:
path = osp.join(root, name)
degree = feat_str.find("deg") >= 0
onehot_maxdeg = re.findall("odeg(\d+)", feat_str)
onehot_maxdeg = int(onehot_maxdeg[0]) if onehot_maxdeg else None
k = re.findall("an{0,1}k(\d+)", feat_str)
k = int(k[0]) if k else 0
groupd = re.findall("groupd(\d+)", feat_str)
groupd = int(groupd[0]) if groupd else 0
remove_edges = re.findall("re(\w+)", feat_str)
remove_edges = remove_edges[0] if remove_edges else 'none'
centrality = feat_str.find("cent") >= 0
coord = feat_str.find("coord") >= 0
pre_transform = FeatureExpander(degree=degree,
onehot_maxdeg=onehot_maxdeg,
AK=k,
centrality=centrality,
remove_edges=remove_edges,
group_degree=groupd).transform
if 'MNIST' in name or 'CIFAR' in name:
if name == 'MNIST_SUPERPIXEL':
train_dataset = MNISTSuperpixels(path,
True,
pre_transform=pre_transform,
transform=T.Cartesian())
test_dataset = MNISTSuperpixels(path,
False,
pre_transform=pre_transform,
transform=T.Cartesian())
else:
train_dataset = ImageDataset(path,
name,
True,
pre_transform=pre_transform,
coord=coord,
processed_file_prefix="data_%s" %
feat_str)
test_dataset = ImageDataset(path,
name,
False,
pre_transform=pre_transform,
coord=coord,
processed_file_prefix="data_%s" %
feat_str)
dataset = (train_dataset, test_dataset)
elif 'QM9' in name:
dataset = QM9Ext(path,
pre_transform=pre_transform,
processed_filename="data_%s.pt" % feat_str)
elif 'ModelNet' in name:
pre_transform = FeatureExpander(
degree=degree,
onehot_maxdeg=onehot_maxdeg,
AK=k,
centrality=centrality,
remove_edges=remove_edges,
group_degree=groupd).cloud_point_transform
train_dataset = ModelNetExT(path,
train=True,
pre_transform=pre_transform,
processed_file_prefix="data_%s" % feat_str)
test_dataset = ModelNetExT(path,
train=True,
pre_transform=pre_transform,
processed_file_prefix="data_%s" % feat_str)
dataset = (train_dataset, test_dataset)
elif 'TOSCA' in name:
# pre_transform = FeatureExpander(
# degree=degree, onehot_maxdeg=onehot_maxdeg, AK=k,
# centrality=centrality, remove_edges=remove_edges,
# group_degree=groupd).cloud_point_transform
dataset = TOSCAEXT(path,
pre_transform=pre_transform,
processed_file_prefix="data_%s" % feat_str)
else:
dataset = TUDatasetExt(path,
name,
pre_transform=pre_transform,
use_node_attr=True,
processed_filename="data_%s.pt" % feat_str)
dataset.data.edge_attr = None
return dataset
def prepare_data(self):
self.train_dataset = MNISTSuperpixels(args.data_fp, True, pre_transform=T.Compose([remove_self_loops, T.Polar()]))
self.test_dataset = MNISTSuperpixels(args.data_fp, False, pre_transform=T.Compose([remove_self_loops, T.Polar()]))
# print(data.num_nodes) # print(data.num_edges) # print(data.num_features) # print(data.is_directed()) # Loading benchmark datasets from torch_geometric.data import InMemoryDataset from torch_geometric.datasets import MNISTSuperpixels, TUDataset, ModelNet dataset = 'Mnist' path = '../data/geometric/MNIST' dataset = MNISTSuperpixels(path, train=True) split = int(len(dataset) * .9) # Randomly permute dataset before splitting dataset = dataset.shuffle() print(dataset.num_classes) print(dataset.num_features) print(dataset[0]) train = dataset[:split] test = dataset[split:] # dataset = ModelNet(root='../data/ModelNet', name='10') # print(len(dataset)) # print(dataset.num_classes)
import os.path as osp
import torch
import torch.nn.functional as F
from torch_geometric.datasets import MNISTSuperpixels
from torch_geometric.data import DataListLoader
from torch_geometric.nn import GCNConv, global_mean_pool, DataParallel
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'MNIST')
dataset = MNISTSuperpixels(path).shuffle()
loader = DataListLoader(dataset, batch_size=512, shuffle=True)
class Net(torch.nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = GCNConv(dataset.num_features, 64)
self.lin = torch.nn.Linear(64, dataset.num_classes)
def forward(self, data):
print('Inside Model: num graphs: {}, device: {}'.format(
data.num_graphs, data.batch.device))
x = self.conv1(data.x, data.edge_index)
x = global_mean_pool(x, data.batch)
return F.log_softmax(self.lin(x), dim=-1)
model = Net()
if torch.cuda.device_count() > 1:
print('Let\'s use', torch.cuda.device_count(), 'GPUs!')
def main(args, kwargs):
dataset = args.data.upper()
if dataset in ['CORA', 'CITESEER', 'PUBMED']:
path = '../data/geometric/' + dataset
print('Using {} dataset'.format(dataset))
dataset = Planetoid(path, dataset, T.NormalizeFeatures())
elif dataset == 'MNIST':
path = '../data/geometric/' + dataset
print('Using {} dataset'.format(dataset))
dataset = MNISTSuperpixels(path, dataset, T.NormalizeFeatures())
else:
sys.exit(
"You must choose one of the 'Planetoid' datasets (CORA, CITESEER, or PUBMED)."
)
data = dataset[0]
# Store the original adjacnecy matrix (for later calculation of edge prediction accuracy)
adj_original = torch.tensor(np.asarray(get_adjacency(data).toarray(),
dtype=np.float32),
requires_grad=False)
channels = 16
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = kwargs[args.model](Encoder(dataset.num_features,
channels)).to(device)
data.train_mask = data.val_mask = data.test_mask = data.y = None
data = model.split_edges(data)
x, edge_index = data.x.to(device), data.edge_index.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
# Create defaultdict of lists to be used for keeping running tally of accuracy, AUC, and AP scores
results = defaultdict(list)
def train():
model.train()
optimizer.zero_grad()
# Produces N * channels vector
z = model.encode(x, edge_index)
'''
TESTING DECODE_INDICES
'''
# pos_decoded = model.decode_indices(z, data.train_pos_edge_index)
# print('z shape: ', z.detach().numpy().shape)
# print('pos edge index shape: ', data.train_pos_edge_index.detach().numpy().shape)
# print('decode_indices_pos: ', pos_decoded.detach().numpy())
# print('decode_indices_pos shape: ', pos_decoded.detach().numpy().shape)
# neg_edge_index = negative_sampling(data.train_pos_edge_index, z.size(0))
# neg_decoded = model.decode_indices(z, neg_edge_index)
# print('decode_indices_neg: ', neg_decoded.detach().numpy())
# print('decode_indices_neg shape: ', neg_decoded.detach().numpy().shape)
'''
END TESTING
'''
'''
Testing Kernel Similarity
'''
# g1 = nx.from_scipy_sparse_matrix(adj_original)
# print(g1)
# from sklearn.metrics.pairwise import rbf_kernel
# adj_reconstructed = model.decode(z, sigmoid=True).detach().numpy()
# adj_reconstructed = (adj_reconstructed > .5).astype(int)
# similarity = kernel_similarity(adj_original, adj_reconstructed, rbf_kernel)
# print(similarity)
'''
END TESTING
'''
if args.loss == 'bce':
loss = model.recon_loss(z, data.train_pos_edge_index)
results['recon_loss'].append(loss)
kl_loss = model.kl_loss()
results['kl'].append(kl_loss)
loss = loss + 0.001 * kl_loss
elif args.loss == 'newbce':
numNodes = len(data['x'])
loss = model.new_recon_loss(z,
edge_index,
num_nodes=numNodes,
num_channels=channels,
adj_original=adj_original)
results['recon_loss'].append(loss)
kl_loss = model.kl_loss()
results['kl'].append(kl_loss)
loss = loss + 0.001 * kl_loss
elif args.loss == 'l2':
loss = model.recon_loss_l2(z, adj_original)
results['recon_loss'].append(loss)
kl_loss = model.kl_loss()
results['kl'].append(kl_loss)
loss = loss + 0.001 * kl_loss
elif args.loss == 'anneal':
# anneal_weight = (1.0 - args.kl_weight) / (len(data.train_pos_edge_index) )
# linear annealing
anneal_rate = epoch / args.kl_warmup if epoch < args.kl_warmup else 1.
kl_weight = min(1.0, anneal_rate)
# Starting annealing
# kl_weight = args.kl_start
# anneal_rate = (1.0 - args.kl_start) / (args.kl_warmup * (len(data)))
# kl_weight = min(1.0, kl_weight + anneal_rate)
print(kl_weight)
loss = model.recon_loss(z, data.train_pos_edge_index)
results['recon_loss'].append(loss)
kl_loss = model.kl_loss() * kl_weight
results['kl'].append(kl_loss)
loss = loss + 0.001 * (kl_loss)
results['loss'].append(loss)
loss.backward()
optimizer.step()
def validate(pos_edge_index, neg_edge_index):
model.eval()
with torch.no_grad():
z = model.encode(x, edge_index)
# accuracy = model.get_accuracy(z, pos_edge_index, neg_edge_index, adj_original)
# accuracy = model.get_accuracy_new(z, adj_original)
# print(accuracy)
auc, ap = model.test(z, pos_edge_index, neg_edge_index)
# return accuracy, auc, ap
return auc, ap
def test(pos_edge_index, neg_edge_index):
model.eval()
with torch.no_grad():
z = model.encode(x, edge_index)
# accuracy = model.get_accuracy(z, pos_edge_index, neg_edge_index, adj_original)
# accuracy = model.get_accuracy_new(z, adj_original)
auc, ap = model.test(z, pos_edge_index, neg_edge_index)
# return accuracy, auc, ap
return auc, ap
for epoch in range(1, args.epochs):
train()
# Run on validation edges
# accuracy, auc, ap = validate(data.val_pos_edge_index, data.val_neg_edge_index)
auc, ap = validate(data.val_pos_edge_index, data.val_neg_edge_index)
# results['acc_val'].append(accuracy)
results['auc_val'].append(auc)
results['ap_val'].append(ap)
# Print AUC and AP on validation data epochs
if epoch % 10 == 0:
# print('Val Epoch : {:03d}, ACC: {:.4f}, AUC: {:.4f}, AP: {:.4f}'.format(epoch, accuracy, auc, ap))
print('Val Epoch : {:03d}, AUC: {:.4f}, AP: {:.4f}'.format(
epoch, auc, ap))
# Evalulate on heldout test edges for every epoch which is a multiple of the test_freq argument
if epoch % args.test_freq == 0:
# accuracy, auc, ap = test(data.test_pos_edge_index, data.test_neg_edge_index)
auc, ap = test(data.test_pos_edge_index, data.test_neg_edge_index)
# results['acc_test'].append(accuracy)
results['auc_test'].append(auc)
results['ap_test'].append(ap)
# print('Test Epoch: {:03d}, ACC: {:.4f}, AUC: {:.4f}, AP: {:.4f}'.format(epoch, accuracy, auc, ap))
print('Test Epoch: {:03d}, AUC: {:.4f}, AP: {:.4f}'.format(
epoch, auc, ap))
# Evaluate on held-out test edges
# auc, ap = test(data.test_pos_edge_index, data.test_neg_edge_index)
# print('Test AUC: {:.4f}, Test AP: {:.4f}'.format(auc, ap))
# Pickle results
if args.save:
if args.notes is None:
modelPath = '../models/' + args.data + '_RESULTS.p'
plotPath = '../figures/geometric/' + args.data + '_RESULTS.png'
lossPath = '../figures/geometric/' + args.data + '_LOSSES.png'
else:
modelPath = '../models/' + args.notes + '_' + args.data + '_RESULTS.p'
plotPath = '../figures/geometric/' + args.notes + '_' + args.data + '_RESULTS.png'
lossPath = '../figures/geometric/' + args.notes + '_' + args.data + '_LOSSES.png'
pkl.dump(results, open(modelPath, 'wb'))
plot_results(pkl.load(open(modelPath, 'rb')),
path=plotPath,
loss=args.loss,
anneal=args.kl_warmup)
plot_losses(pkl.load(open(modelPath, 'rb')),
path=lossPath,
loss=args.loss,
anneal=args.kl_warmup)
def load_dataset(args):
# automatic data loading and splitting
transform = add_zeros if args.dataset == 'ogbg-ppa' else None
cls_criterion = get_loss_function(args.dataset)
idx2word_mapper = None
if args.dataset == 'mnist':
train_data = MNISTSuperpixels(root='dataset',
train=True,
transform=T.Polar())
dataset = train_data
dataset.name = 'mnist'
dataset.eval_metric = 'acc'
validation_data = []
test_data = MNISTSuperpixels(root='dataset',
train=False,
transform=T.Polar())
train_data = list(train_data)
test_data = list(test_data)
elif args.dataset == 'QM9':
# Contains 19 targets. Use only the first 12 (0-11)
QM9_VALIDATION_START = 110000
QM9_VALIDATION_END = 120000
dataset = QM9(root='dataset',
transform=ExtractTargetTransform(args.target)).shuffle()
dataset.name = 'QM9'
dataset.eval_metric = 'mae'
train_data = dataset[:QM9_VALIDATION_START]
validation_data = dataset[QM9_VALIDATION_START:QM9_VALIDATION_END]
test_data = dataset[QM9_VALIDATION_END:]
train_data = list(train_data)
validation_data = list(validation_data)
test_data = list(test_data)
elif args.dataset == 'zinc':
train_data = ZINC(root='dataset', subset=True, split='train')
dataset = train_data
dataset.name = 'zinc'
validation_data = ZINC(root='dataset', subset=True, split='val')
test_data = ZINC(root='dataset', subset=True, split='test')
dataset.eval_metric = 'mae'
train_data = list(train_data)
validation_data = list(validation_data)
test_data = list(test_data)
elif args.dataset in [
'ogbg-molhiv', 'ogbg-molpcba', 'ogbg-ppa', 'ogbg-code2'
]:
dataset = PygGraphPropPredDataset(name=args.dataset,
transform=transform)
if args.dataset == 'obgb-code2':
seq_len_list = np.array([len(seq) for seq in dataset.data.y])
max_seq_len = args.max_seq_len
num_less_or_equal_to_max = np.sum(
seq_len_list <= args.max_seq_len) / len(seq_len_list)
print(
f'Target sequence less or equal to {max_seq_len} is {num_less_or_equal_to_max}%.'
)
split_idx = dataset.get_idx_split()
# The following is only used in the evaluation of the ogbg-code classifier.
if args.dataset == 'ogbg-code2':
vocab2idx, idx2vocab = get_vocab_mapping(
[dataset.data.y[i] for i in split_idx['train']],
args.num_vocab)
# specific transformations for the ogbg-code dataset
dataset.transform = transforms.Compose([
augment_edge,
lambda data: encode_y_to_arr(data, vocab2idx, args.max_seq_len)
])
idx2word_mapper = partial(decode_arr_to_seq, idx2vocab=idx2vocab)
train_data = list(dataset[split_idx["train"]])
validation_data = list(dataset[split_idx["valid"]])
test_data = list(dataset[split_idx["test"]])
return dataset, train_data, validation_data, test_data, cls_criterion, idx2word_mapper
def TrainingNet(dataset,modelName,params,num_pre_epochs,num_epochs,NumCutoff,optimizerName,LinkPredictionMethod,MonteSize,savepath):
Batch_size=params[0]
VectorPairs=params[4]
StartTopoCoeffi=params[5]
WeightCorrectionCoeffi=params[6]
interval=params[7]
root='/git/data/GraphData/'+dataset
TestAccs=[]
for Monte_iter in range(MonteSize):
# Data
NewNetworkSizeAdjust=[]
WeightsDynamicsEvolution=[]
trainValRatio=[0.2,0.4]
# model
if dataset=='Cora' or dataset =='Citeseer' or dataset =='Pubmed':
datasetroot= Planetoid(root=root, name=dataset,transform =T.NormalizeFeatures()).shuffle()
trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True)
""" train_mask, val_mask,test_mask=DataSampler(trainValRatio,datasetroot.data.num_nodes)
DataMask={}
DataMask['train_mask']=train_mask
DataMask['val_mask']=val_mask
DataMask['test_mask']=test_mask
trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True)"""
num_features=datasetroot.num_features
num_classes=datasetroot.num_classes
criterion = nn.CrossEntropyLoss()
elif dataset =="CoraFull":
datasetroot = CoraFull(root=root,transform =T.NormalizeFeatures()).shuffle()
"""train_mask, val_mask,test_mask=DataSampler(trainValRatio,datasetroot.data.num_nodes)
DataMask={}
DataMask['train_mask']=train_mask
DataMask['val_mask']=val_mask
DataMask['test_mask']=test_mask"""
criterion = nn.CrossEntropyLoss()
num_features=datasetroot.num_features
num_classes=datasetroot.num_classes
trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=False)
elif dataset=='ENZYMES' or dataset=='MUTAG':
datasetroot=TUDataset(root,name=dataset,use_node_attr=True)
trainloader = DataLoader(datasetroot, batch_size=Batch_size, shuffle=True)
num_features=datasetroot.num_features
num_classes=datasetroot.num_classes
elif dataset =="PPI":
train_dataset = PPI(root, split='train')
val_dataset = PPI(root, split='val')
test_dataset = PPI(root, split='test')
trainloader = DataListLoader(train_dataset, batch_size=Batch_size, shuffle=True)
valloader = DataListLoader(val_dataset, batch_size=100, shuffle=False)
testloader = DataListLoader(test_dataset, batch_size=100, shuffle=False)
num_classes=train_dataset.num_classes
num_features=train_dataset.num_features
criterion = torch.nn.BCEWithLogitsLoss()
elif dataset =="Reddit":
datasetroot=Reddit(root)
trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True)
testloader = DataListLoader(datasetroot, batch_size=2, shuffle=False)
elif dataset=="Amazon":
datasetroot=Amazon(root, "Photo", transform=None, pre_transform=None)
trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True)
testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False)
elif dataset=='MNIST':
datasetroot = MNISTSuperpixels(root=root, transform=T.Cartesian())
trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True)
testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False)
elif dataset=='CIFAR10':
pass
else:
raise Exception("Input wrong datatset!!")
width=ContractionLayerCoefficients(num_features,*params[1:3])
net =ChooseModel(modelName,num_features,num_classes,width)
FileName="{}-{}-param_{}_{}_{}_{}-monte_{}".format(dataset,modelName,interval,WeightCorrectionCoeffi,StartTopoCoeffi,VectorPairs,Monte_iter)
print('Let\'s use', torch.cuda.device_count(), 'GPUs!')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
criterion = criterion.to(device)
net = DataParallel(net)
net = net.to(device)
optimizer = getattr(optim,optimizerName)(net.parameters(), lr=params[3], momentum=0.9, weight_decay=5e-4)
model_to_save='./checkpoint/{}-{}-param_{}_{}_{}_{}-ckpt.pth'.format(dataset,modelName,params[0],params[1],params[5],params[4])
if resume=="True" and os.path.exists(model_to_save):
[net,optimizer,TrainConvergence,TestConvergence,Acc]=ResumeModel(net,optimizer,model_to_save)
start_epoch=len(TrainConvergence)
else:
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
#cudnn.benchmark = True
logging('dataset:{}, Batch size: {}, Number of layers:{} ConCoeff: {}, LR:{}, MonteSize:{}'.format(dataset, params[0], params[1],params[2],params[3],Monte_iter))
mark="{}{}Convergence/DiagElement-{}".format(savepath,dataset,FileName)
markweights="{}{}Convergence/WeightChanges-{}".format(savepath,dataset,FileName)
PreTrainConvergence,PreTestConvergence,PreAcc=TrainPart(start_epoch,num_pre_epochs,num_classes, trainloader,net,optimizer,criterion,NumCutoff,LinkPredictionMethod,VectorPairs,WeightCorrectionCoeffi,StartTopoCoeffi,mark,markweights,model_to_save,False)
print('dataset: {}, model name:{}, epoch:{},Pre-train error:{}; Pre-test error:{}; test acc:{}'.format(dataset,modelName,num_pre_epochs,PreTrainConvergence[-1],PreTestConvergence[-1],PreAcc))
NewNetworksize=RetainNetworkSize(net,params[2])
OptimizedNet=ChooseModel(modelName,num_features,num_classes,NewNetworksize[0:-1])
NewNetworksize.insert(0,num_features)
NewNetworkSizeAdjust.append(NewNetworksize[0:-1])
print(NewNetworkSizeAdjust)
#OptimizedNet.apply(init_weights)
OptimizedNet = DataParallel(OptimizedNet)
OptimizedNet = OptimizedNet.to(device)
cudnn.benchmark = True
# Begin Pre training
if optimizerName =="SGD":
optimizerNew = getattr(optim,optimizerName)(OptimizedNet.parameters(), lr=params[3], momentum=0.9, weight_decay=5e-4)
elif optimizerName =="Adam":
optimizerNew = getattr(optim,optimizerName)(OptimizedNet.parameters(), lr=params[3], betas=(0.9, 0.999), eps=1e-08, weight_decay=5e-4, amsgrad=False)
TrainConvergence,TestConvergence,TestAcc=TrainPart(start_epoch,num_epochs,datasetroot.num_classes, trainloader,OptimizedNet,optimizerNew,criterion,
NumCutoff,LinkPredictionMethod,VectorPairs,WeightCorrectionCoeffi,StartTopoCoeffi,mark,markweights,model_to_save,True)
np.save("{}/{}Convergence/AlgebraicConectivityTrainConvergence-{}".format(savepath,dataset,FileName),TrainConvergence)
np.save("{}/{}Convergence/AlgebraicConectivityTestConvergence-{}".format(savepath,dataset,FileName),TestConvergence)
#np.save("{}/{}Convergence/NewNetworkSizeAdjust-{}".format(savepath,dataset,FileName),NewNetworkSizeAdjust)
#torch.cuda.empty_cache()
print('dataset: {}, model name:{}, resized network size: {}, the train error: {},test error: {}, test acc:{}\n'.format(dataset,modelName,NewNetworksize[0:-1],num_epochs,TrainConvergence[-1],TestConvergence[-1],TestAcc))
np.save("{}/{}Convergence/AlgebraicConectivityMeanTestAccs-{}".format(savepath,dataset,FileName),TestAccs.append(TestAcc))
TestAccs.append(TestAcc)
print_nvidia_useage()
def GCN(dataset, params, Epochs, MonteSize, width, lr, savepath):
Batch_size = int(params[0])
for Monte_iter in range(MonteSize):
# Data
best_loss = float('inf') # best test loss
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
TrainConvergence = []
TestConvergence = []
# model
root = '/data/GraphData/' + dataset
if dataset == 'Cora':
model_name = "GCN3"
datasetroot = Planetoid(root=root, name=dataset).shuffle()
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
testloader = DataListLoader(datasetroot,
batch_size=100,
shuffle=False)
model_to_save = './checkpoint/{}-{}-param_{}_{}-Mon_{}-ckpt.pth'.format(
dataset, model_name, params[0], params[1], Monte_iter)
if resume and os.path.exists(model_to_save):
[net, TrainConvergence, TestConvergence,
start_epoch] = ResumeModel(model_to_save)
if start_epoch >= Epochs - 1:
continue
else:
net = Net(datasetroot, width)
elif dataset == 'ENZYMES' or dataset == 'MUTAG':
model_name = "topk_pool_Net"
root = '/data/GraphData' + dataset
datasetroot = TUDataset(root, name=dataset)
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
testloader = DataListLoader(datasetroot,
batch_size=100,
shuffle=False)
model_to_save = './checkpoint/{}-{}-param_{}_{}-Mon_{}-ckpt.pth'.format(
dataset, model_name, params[0], params[1], Monte_iter)
if resume and os.path.exists(model_to_save):
[net, TrainConvergence, TestConvergence,
start_epoch] = ResumeModel(model_to_save)
if start_epoch >= Epochs - 1:
continue
else:
net = topk_pool_Net(datasetroot, width)
elif dataset == 'MNIST':
datasetroot = MNISTSuperpixels(root='/data/GraphData/' + dataset,
transform=T.Cartesian()).shuffle()
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
testloader = DataListLoader(datasetroot,
batch_size=100,
shuffle=False)
model_name = 'SPlineNet'
model_to_save = './checkpoint/{}-{}-param_{}_{}-Mon_{}-ckpt.pth'.format(
dataset, model_name, params[0], params[1], Monte_iter)
if resume and os.path.exists(model_to_save):
[net, TrainConvergence, TestConvergence,
start_epoch] = ResumeModel(model_to_save)
if start_epoch >= Epochs - 1:
continue
else:
#net=Net(datasetroot,width)
net = SPlineNet(datasetroot, width)
elif dataset == 'CIFAR10':
if resume and os.path.exists(model_to_save):
[net, TrainConvergence, TestConvergence,
start_epoch] = ResumeModel(model_to_save)
if start_epoch >= Epochs - 1:
continue
else:
net = getattr(CIFAR10_resnet, 'Resnet20_CIFAR10')(params[1])
else:
raise Exception(
"The dataset is:{}, it isn't existed.".format(dataset))
print('Let\'s use', torch.cuda.device_count(), 'GPUs!')
torch.cuda.is_available()
net = DataParallel(net)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
net = net.to(device)
#cudnn.benchmark = True
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=5e-4)
for epoch in range(start_epoch, start_epoch + Epochs):
if epoch < Epochs:
logging(
'Batch size: {},ConCoeff: {},MonteSize:{},epoch:{}'.format(
params[0], params[1], Monte_iter, epoch))
TrainLoss = train(trainloader, net, optimizer, criterion)
TrainConvergence.append(statistics.mean(TrainLoss))
TestConvergence.append(
statistics.mean(test(testloader, net, criterion)))
else:
break
if TestConvergence[epoch] < best_loss:
logging('Saving..')
state = {
'net': net.module,
'TrainConvergence': TrainConvergence,
'TestConvergence': TestConvergence,
'epoch': epoch,
}
if not os.path.isdir('checkpoint'):
os.mkdir('checkpoint')
torch.save(state, model_to_save)
best_loss = TestConvergence[epoch]
if not os.path.exists('./%s' % model_name):
os.makedirs('./%s' % model_name)
torch.save(
net.module.state_dict(),
'./%s/%s_%s_%s_%s_%s_pretrain.pth' %
(model_name, dataset, model_name, params[0], params[1],
Epochs))
else:
pass
## save recurrence plots
if epoch % 20 == 0:
save_recurrencePlots_file = "../Results/RecurrencePlots/RecurrencePlots_{}_{}_BatchSize{}_ConCoeffi{}_epoch{}.png".format(
dataset, model_name, params[0], params[1], epoch)
save_recurrencePlots(net, save_recurrencePlots_file)
FileName = "{}-{}-param_{}_{}-monte_{}".format(dataset, model_name,
params[0], params[1],
Monte_iter)
np.save(savepath + 'TrainConvergence-' + FileName, TrainConvergence)
np.save(savepath + 'TestConvergence-' + FileName, TestConvergence)
torch.cuda.empty_cache()
print_nvidia_useage()
if return_output == True:
return TestConvergence[-1], net.module.fc.weight
else:
pass
def main(args):
args = init(args)
def pf(data):
return data.y == args.num
pre_filter = pf if args.num != -1 else None
print("loading data")
if (args.sparse_mnist):
X = MNISTGraphDataset(args.dataset_path,
args.num_hits,
train=args.train,
num=args.num)
X_loaded = DataLoader(X,
shuffle=True,
batch_size=args.batch_size,
pin_memory=True)
else:
if (args.gcnn):
X = MNISTSuperpixels(args.dir_path,
train=args.train,
pre_transform=T.Cartesian(),
pre_filter=pre_filter)
X_loaded = tgDataLoader(X,
shuffle=True,
batch_size=args.batch_size)
else:
X = SuperpixelsDataset(args.dataset_path,
args.num_hits,
train=args.train,
num=args.num)
X_loaded = DataLoader(X,
shuffle=True,
batch_size=args.batch_size,
pin_memory=True)
print("loaded data")
# model
if (args.load_model):
G = torch.load(args.model_path + args.name + "/G_" +
str(args.start_epoch) + ".pt",
map_location=args.device)
D = torch.load(args.model_path + args.name + "/D_" +
str(args.start_epoch) + ".pt",
map_location=args.device)
else:
# G = Graph_Generator(args.node_feat_size, args.fe_hidden_size, args.fe_out_size, args.fn_hidden_size, args.fn_num_layers, args.mp_iters_gen, args.num_hits, args.gen_dropout, args.leaky_relu_alpha, hidden_node_size=args.hidden_node_size, int_diffs=args.int_diffs, pos_diffs=args.pos_diffs, gru=args.gru, batch_norm=args.batch_norm, device=device).to(args.device)
if (args.gcnn):
G = GaussianGenerator(args=deepcopy(args)).to(args.device)
D = MoNet(args=deepcopy(args)).to(args.device)
# D = Gaussian_Discriminator(args.node_feat_size, args.fe_hidden_size, args.fe_out_size, args.mp_hidden_size, args.mp_num_layers, args.num_iters, args.num_hits, args.dropout, args.leaky_relu_alpha, kernel_size=args.kernel_size, hidden_node_size=args.hidden_node_size, int_diffs=args.int_diffs, gru=GRU, batch_norm=args.batch_norm, device=device).to(args.device)
else:
# D = Graph_Discriminator(args.node_feat_size, args.fe_hidden_size, args.fe_out_size, args.fn_hidden_size, args.fn_num_layers, args.mp_iters_disc, args.num_hits, args.disc_dropout, args.leaky_relu_alpha, hidden_node_size=args.hidden_node_size, wgan=args.wgan, int_diffs=args.int_diffs, pos_diffs=args.pos_diffs, gru=args.gru, batch_norm=args.batch_norm, device=device).to(args.device)
print("Generator")
G = Graph_GAN(gen=True, args=deepcopy(args)).to(args.device)
print("Discriminator")
D = Graph_GAN(gen=False, args=deepcopy(args)).to(args.device)
print("Models loaded")
# optimizer
if args.spectral_norm_gen:
G_params = filter(lambda p: p.requires_grad, G.parameters())
else:
G_params = G.parameters()
if args.spectral_norm_gen:
D_params = filter(lambda p: p.requires_grad, D.parameters())
else:
D_params = D.parameters()
if (args.optimizer == 'rmsprop'):
G_optimizer = optim.RMSprop(G_params, lr=args.lr_gen)
D_optimizer = optim.RMSprop(D_params, lr=args.lr_disc)
elif (args.optimizer == 'adadelta'):
G_optimizer = optim.Adadelta(G_params, lr=args.lr_gen)
D_optimizer = optim.Adadelta(D_params, lr=args.lr_disc)
elif (args.optimizer == 'acgd'):
optimizer = ACGD(max_params=G_params,
min_params=D_params,
lr_max=args.lr_gen,
lr_min=args.lr_disc,
device=args.device)
elif (args.optimizer == 'adam' or args.optimizer == 'None'):
G_optimizer = optim.Adam(G_params,
lr=args.lr_gen,
weight_decay=5e-4,
betas=(args.beta1, args.beta2))
D_optimizer = optim.Adam(D_params,
lr=args.lr_disc,
weight_decay=5e-4,
betas=(args.beta1, args.beta2))
if (args.load_model):
try:
if (not args.optimizer == 'acgd'):
G_optimizer.load_state_dict(
torch.load(args.model_path + args.name + "/G_optim_" +
str(args.start_epoch) + ".pt",
map_location=args.device))
D_optimizer.load_state_dict(
torch.load(args.model_path + args.name + "/D_optim_" +
str(args.start_epoch) + ".pt",
map_location=args.device))
else:
optimizer.load_state_dict(
torch.load(args.model_path + args.name + "/optim_" +
str(args.start_epoch) + ".pt",
map_location=args.device))
except:
print("Error loading optimizer")
print("optimizers loaded")
if args.fid: C, mu2, sigma2 = evaluation.load(args, X_loaded)
normal_dist = Normal(
torch.tensor(0.).to(args.device),
torch.tensor(args.sd).to(args.device))
lns = args.latent_node_size if args.latent_node_size else args.hidden_node_size
args.noise_file_name = "num_samples_" + str(
args.num_samples) + "_num_nodes_" + str(
args.num_hits) + "_latent_node_size_" + str(lns) + "_sd_" + str(
args.sd) + ".pt"
if args.gcnn: args.noise_file_name = "gcnn_" + args.noise_file_name
noise_file_names = listdir(args.noise_path)
if args.noise_file_name not in noise_file_names:
if (args.gcnn):
torch.save(
normal_dist.sample(
(args.num_samples * 5, 2 + args.channels[0])),
args.noise_path + args.noise_file_name)
else:
torch.save(
normal_dist.sample((args.num_samples, args.num_hits, lns)),
args.noise_path + args.noise_file_name)
losses = {}
if (args.load_model):
try:
losses['D'] = np.loadtxt(args.losses_path + args.name + "/" +
"D.txt").tolist()[:args.start_epoch]
losses['Dr'] = np.loadtxt(args.losses_path + args.name + "/" +
"Dr.txt").tolist()[:args.start_epoch]
losses['Df'] = np.loadtxt(args.losses_path + args.name + "/" +
"Df.txt").tolist()[:args.start_epoch]
losses['G'] = np.loadtxt(args.losses_path + args.name + "/" +
"G.txt").tolist()[:args.start_epoch]
if args.fid:
losses['fid'] = np.loadtxt(
args.losses_path + args.name + "/" +
"fid.txt").tolist()[:args.start_epoch]
if (args.gp):
losses['gp'] = np.loadtxt(args.losses_path + args.name + "/" +
"gp.txt").tolist()[:args.start_epoch]
except:
print("couldn't load losses")
losses['D'] = []
losses['Dr'] = []
losses['Df'] = []
losses['G'] = []
if args.fid: losses['fid'] = []
if (args.gp): losses['gp'] = []
else:
losses['D'] = []
losses['Dr'] = []
losses['Df'] = []
losses['G'] = []
if args.fid: losses['fid'] = []
if (args.gp): losses['gp'] = []
Y_real = torch.ones(args.batch_size, 1).to(args.device)
Y_fake = torch.zeros(args.batch_size, 1).to(args.device)
def train_D(data, gen_data=None, unrolled=False):
if args.debug: print("dtrain")
D.train()
D_optimizer.zero_grad()
run_batch_size = data.shape[0] if not args.gcnn else data.y.shape[0]
if gen_data is None:
gen_data = utils.gen(args, G, normal_dist, run_batch_size)
if (args.gcnn):
gen_data = utils.convert_to_batch(args, gen_data,
run_batch_size)
if args.augment:
p = args.aug_prob if not args.adaptive_prob else losses['p'][-1]
data = augment.augment(args, data, p)
gen_data = augment.augment(args, gen_data, p)
D_real_output = D(data.clone())
D_fake_output = D(gen_data)
D_loss, D_loss_items = utils.calc_D_loss(args, D, data, gen_data,
D_real_output, D_fake_output,
run_batch_size, Y_real,
Y_fake)
D_loss.backward(create_graph=unrolled)
D_optimizer.step()
return D_loss_items
def train_G(data):
if args.debug: print("gtrain")
G.train()
G_optimizer.zero_grad()
gen_data = utils.gen(args, G, normal_dist, args.batch_size)
if (args.gcnn):
gen_data = utils.convert_to_batch(args, gen_data, args.batch_size)
if args.augment:
p = args.aug_prob if not args.adaptive_prob else losses['p'][-1]
gen_data = augment.augment(args, gen_data, p)
if (args.unrolled_steps > 0):
D_backup = deepcopy(D)
for i in range(args.unrolled_steps - 1):
train_D(data, gen_data=gen_data, unrolled=True)
D_fake_output = D(gen_data)
G_loss = utils.calc_G_loss(args, D_fake_output, Y_real)
G_loss.backward()
G_optimizer.step()
if (args.unrolled_steps > 0):
D.load(D_backup)
return G_loss.item()
def train_acgd(data):
if args.debug: print("acgd train")
D.train()
G.train()
optimizer.zero_grad()
run_batch_size = data.shape[0] if not args.gcnn else data.y.shape[0]
gen_data = utils.gen(args, G, normal_dist, run_batch_size)
if (args.gcnn):
gen_data = utils.convert_to_batch(args, gen_data, run_batch_size)
if args.augment:
p = args.aug_prob if not args.adaptive_prob else losses['p'][-1]
data = utils.rand_translate(args, data, p)
gen_data = utils.rand_translate(args, gen_data, p)
D_real_output = D(data.clone())
D_fake_output = D(gen_data)
D_loss, D_loss_items = utils.calc_D_loss(args, D, data, gen_data,
D_real_output, D_fake_output,
run_batch_size)
optimizer.step(loss=D_loss)
G.eval()
with torch.no_grad():
G_loss = utils.calc_G_loss(args, D_fake_output)
return D_loss_items, G_loss.item()
def train():
k = 0
temp_ng = args.num_gen
if (args.fid):
losses['fid'].append(
evaluation.get_fid(args, C, G, normal_dist, mu2, sigma2))
if (args.save_zero):
save_outputs.save_sample_outputs(args, D, G, normal_dist,
args.name, 0, losses)
for i in range(args.start_epoch, args.num_epochs):
print("Epoch %d %s" % ((i + 1), args.name))
Dr_loss = 0
Df_loss = 0
G_loss = 0
D_loss = 0
gp_loss = 0
lenX = len(X_loaded)
for batch_ndx, data in tqdm(enumerate(X_loaded), total=lenX):
data = data.to(args.device)
if (args.gcnn):
data.pos = (data.pos - 14) / 28
row, col = data.edge_index
data.edge_attr = (data.pos[col] -
data.pos[row]) / (2 * args.cutoff) + 0.5
if (not args.optimizer == 'acgd'):
if (args.num_critic > 1):
D_loss_items = train_D(data)
D_loss += D_loss_items['D']
Dr_loss += D_loss_items['Dr']
Df_loss += D_loss_items['Df']
if (args.gp): gp_loss += D_loss_items['gp']
if ((batch_ndx - 1) % args.num_critic == 0):
G_loss += train_G(data)
else:
if (batch_ndx == 0
or (batch_ndx - 1) % args.num_gen == 0):
D_loss_items = train_D(data)
D_loss += D_loss_items['D']
Dr_loss += D_loss_items['Dr']
Df_loss += D_loss_items['Df']
if (args.gp): gp_loss += D_loss_items['gp']
G_loss += train_G(data)
else:
D_loss_items, G_loss_item = train_acgd(data)
D_loss += D_loss_items['D']
Dr_loss += D_loss_items['Dr']
Df_loss += D_loss_items['Df']
G_loss += G_loss_item
if args.bottleneck:
if (batch_ndx == 10):
return
losses['D'].append(D_loss / (lenX / args.num_gen))
losses['Dr'].append(Dr_loss / (lenX / args.num_gen))
losses['Df'].append(Df_loss / (lenX / args.num_gen))
losses['G'].append(G_loss / (lenX / args.num_critic))
if (args.gp): losses['gp'].append(gp_loss / (lenX / args.num_gen))
print("d loss: " + str(losses['D'][-1]))
print("g loss: " + str(losses['G'][-1]))
print("dr loss: " + str(losses['Dr'][-1]))
print("df loss: " + str(losses['Df'][-1]))
if (args.gp): print("gp loss: " + str(losses['gp'][-1]))
gloss = losses['G'][-1]
drloss = losses['Dr'][-1]
dfloss = losses['Df'][-1]
dloss = (drloss + dfloss) / 2
if (args.bgm):
if (i > 20 and gloss > dloss + args.bag):
print("num gen upping to 10")
args.num_gen = 10
else:
print("num gen normal")
args.num_gen = temp_ng
elif (args.gom):
if (i > 20 and gloss > dloss + args.bag):
print("G loss too high - training G only")
j = 0
print("starting g loss: " + str(gloss))
print("starting d loss: " + str(dloss))
while (gloss > dloss + args.bag * 0.5):
print(j)
gloss = 0
for l in tqdm(range(lenX)):
gloss += train_G()
gloss /= lenX
print("g loss: " + str(gloss))
print("d loss: " + str(dloss))
losses['D'].append(dloss * 2)
losses['Dr'].append(drloss)
losses['Df'].append(dfloss)
losses['G'].append(gloss)
if (j % 5 == 0):
save_outputs.save_sample_outputs(args,
D,
G,
normal_dist,
args.name,
i + 1,
losses,
k=k,
j=j)
j += 1
k += 1
elif (args.rd):
if (i > 20 and gloss > dloss + args.bag):
print("gloss too high, resetting D params")
D.reset_params()
if ((i + 1) % 5 == 0):
optimizers = optimizer if args.optimizer == 'acgd' else (
D_optimizer, G_optimizer)
save_outputs.save_models(args, D, G, optimizers, args.name,
i + 1)
if (args.fid and (i + 1) % 1 == 0):
losses['fid'].append(
evaluation.get_fid(args, C, G, normal_dist, mu2, sigma2))
if ((i + 1) % 5 == 0):
save_outputs.save_sample_outputs(args, D, G, normal_dist,
args.name, i + 1, losses)
train()
def main():
dataset = MNISTSuperpixels(root='~/MNISTSuperpixels', train=True)
dataset_test = MNISTSuperpixels(root='~/MNISTSuperpixels', train=False)
class Net(torch.nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = GCNConv(dataset.num_node_features, 32)
self.conv2 = GCNConv(32, 16)
self.conv3 = GCNConv(16, dataset.num_classes)
#self.conv1 = SAGEConv(dataset.num_node_features, 32)
#self.conv2 = SAGEConv(32, 16)
#self.conv3 = SAGEConv(16, dataset.num_classes)
def forward(self, data):
x, edge_index, batch = data.x, data.edge_index, data.batch
x = self.conv1(x, edge_index)
x = F.relu(x)
#x = F.dropout(x, training=self.training)
x = self.conv2(x, edge_index)
x = F.relu(x)
x = self.conv3(x, edge_index)
x = torch_geometric.nn.global_max_pool(x, batch)
return F.log_softmax(x, dim=1)
loader = DataLoader(dataset, batch_size=32, shuffle=True)
loader_test = DataLoader(dataset_test, batch_size=1, shuffle=False)
for batch in loader:
batch
print(batch)
print(batch.y[0])
print(np.array(batch.y))
break
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net().to(device)
#data = dataset[0].to(device)
optimizer = torch.optim.Adam(model.parameters(),
lr=0.001,
weight_decay=5e-4)
model.train()
n_train_correct = 0
n_train_count = 0
for epoch in range(200):
for batch in loader:
dat = batch.to(device)
optimizer.zero_grad()
out = model(dat)
loss = F.nll_loss(out, dat.y)
loss.backward()
optimizer.step()
n_train_correct += (torch.max(out, 1)[1].view(
batch.y.size()) == batch.y).sum().item()
n_train_count += 32 #batch.y.size()
dev_acc = 100. * n_train_correct / n_train_count
print('Train accuracy: ', dev_acc)
# calculate accuracy on validation set
n_dev_correct = 0
n_count = 0
with torch.no_grad():
for batch in loader_test:
data = batch.to(device)
answer = model(data)
n_dev_correct += (torch.max(answer, 1)[1].view(
batch.y.size()) == batch.y).sum().item()
n_count += 1
dev_acc = 100. * n_dev_correct / n_count
print('Test accuracy: ', dev_acc)
import os.path as osp
import torch
import torch.nn.functional as F
from torch_geometric.datasets import MNISTSuperpixels
from torch_geometric.data import DataListLoader
import torch_geometric.transforms as T
from torch_geometric.nn import SplineConv, global_mean_pool, DataParallel
dataset = MNISTSuperpixels("/git/data/GraphData/MNIST",
transform=T.Cartesian()).shuffle()
loader = DataListLoader(dataset, batch_size=1024, shuffle=True)
class Net(torch.nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = SplineConv(dataset.num_features, 32, dim=2, kernel_size=5)
self.conv2 = SplineConv(32, 64, dim=2, kernel_size=5)
self.lin1 = torch.nn.Linear(64, 128)
self.lin2 = torch.nn.Linear(128, dataset.num_classes)
def forward(self, data):
print('Inside Model: num graphs: {}, device: {}'.format(
data.num_graphs, data.batch.device))
x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr
for i in range(3):
x = F.elu(self.conv1(x, edge_index, edge_attr))
x = F.elu(self.conv2(x, edge_index, edge_attr))
x = global_mean_pool(x, data.batch)
# get the current node index of G
idx = list(G.nodes())
# set the node sizes using node features
size = x[idx] * 500 + 200
# set the node colors using node features
color = []
for i in idx:
grey = x[i]
if grey == 0:
color.append('skyblue')
else:
color.append('red')
nx.draw(G, with_labels=True, node_size=size, node_color=color, pos=pos)
plt.title("MNIST Superpixel")
if __name__ == '__main__':
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'MNIST')
image_dataset = MNIST(root=path, download=True)
graph_dataset = MNISTSuperpixels(root=path)
example = 25
image, label = image_dataset[example]
data = graph_dataset[example]
print("Image Label:", label)
visualize(image, data)
def TrainingNet(dataset, modelName, params, num_pre_epochs, num_epochs,
NumCutoff, optimizerName, MonteSize, savepath):
Batch_size = int(params[0])
root = '/git/data/GraphData/' + dataset
TestAccs = []
for Monte_iter in range(MonteSize):
# Data
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
NewNetworkSizeAdjust = []
WeightsDynamicsEvolution = []
# model
if dataset == 'Cora' or dataset == 'Citeseer' or dataset == 'Pubmed':
datasetroot = Planetoid(root=root,
name=dataset,
transform=T.NormalizeFeatures()).shuffle()
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
[net,
model_to_save] = ModelAndSave(dataset, modelName, datasetroot,
params, num_epochs)
criterion = nn.CrossEntropyLoss()
elif dataset == "CoraFull":
datasetroot = CoraFull(root=root,
transform=T.NormalizeFeatures()).shuffle()
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
[net,
model_to_save] = ModelAndSave(dataset, modelName, datasetroot,
params, num_epochs)
elif dataset == "Amazon":
datasetroot = Amazon(root,
"Photo",
transform=None,
pre_transform=None)
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
testloader = DataListLoader(datasetroot,
batch_size=100,
shuffle=False)
[net,
model_to_save] = ModelAndSave(dataset, modelName, datasetroot,
params, num_epochs)
elif dataset == 'ENZYMES' or dataset == 'MUTAG':
datasetroot = TUDataset(root, name=dataset, use_node_attr=True)
Num = len(datasetroot) // 10
global train_dataset, test_dataset
train_dataset = datasetroot[:Num]
test_dataset = datasetroot[Num:]
trainloader = DataLoader(train_dataset, batch_size=Batch_size)
testloader = DataLoader(test_dataset, batch_size=60)
[net,
model_to_save] = ModelAndSave(dataset, modelName, datasetroot,
params, num_epochs)
elif dataset == "PPI":
train_dataset = PPI(root, split='train')
test_dataset = PPI(root, split='test')
trainloader = DataLoader(train_dataset,
batch_size=Batch_size,
shuffle=True)
testloader = DataLoader(test_dataset, batch_size=1, shuffle=False)
[net,
model_to_save] = ModelAndSave(dataset, modelName, train_dataset,
params, num_epochs)
criterion = torch.nn.BCEWithLogitsLoss()
elif dataset == "Reddit":
datasetroot = Reddit(root)
trainloader = DataListLoader(datasetroot,
batch_size=1,
shuffle=True)
testloader = DataListLoader(datasetroot,
batch_size=2,
shuffle=False)
[net,
model_to_save] = ModelAndSave(dataset, modelName, datasetroot,
params, num_epochs)
criterion = torch.nn.BCEWithLogitsLoss()
elif dataset == 'MNIST':
datasetroot = MNISTSuperpixels(root=root, transform=T.Cartesian())
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
testloader = DataListLoader(datasetroot,
batch_size=100,
shuffle=False)
[net,
model_to_save] = ModelAndSave(dataset, modelName, datasetroot,
params, num_epochs)
elif dataset == 'CIFAR10':
pass
else:
raise Exception("Input wrong datatset!!")
FileName = "{}-{}-param_{}_{}_{}_{}-monte_{}".format(
dataset, modelName, params[0], params[1], params[2], params[3],
Monte_iter)
print('Let\'s use', torch.cuda.device_count(), 'GPUs!')
global device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
optimizer = optim.Adam(net.parameters(),
lr=params[3],
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0,
amsgrad=False)
criterion = nn.CrossEntropyLoss()
net = net.to(device)
#cudnn.benchmark = True
logging(
'dataset:{}, Batch size: {}, Number of layers:{} ConCoeff: {}, LR:{}, MonteSize:{}'
.format(dataset, params[0], params[1], params[2], params[3],
Monte_iter))
mark = "{}/{}Convergence/DiagElement-{}".format(
savepath, dataset, FileName)
PreTrainConvergence, PreTestConvergence, PreTestAcc = TrainPart(
modelName, datasetroot, start_epoch, num_pre_epochs, trainloader,
testloader, net, optimizer, criterion, NumCutoff, mark, False,
model_to_save)
print(
'dataset: {}, model name: {}, Number epoches: {}, Pre-train error is: {}, Pre-test error is: {}, test acc is {}'
.format(dataset, modelName, num_pre_epochs,
PreTrainConvergence[-1], PreTestConvergence[-1],
PreTestAcc[-1]))
NewNetworksize, NewNetworkWeight = RetainNetworkSize(net,
params[2])[0:2]
NetworkInfo = [NewNetworksize[0:-1], NewNetworkWeight]
OptimizedNet = ChooseModel(modelName, datasetroot, NetworkInfo)
NewNetworksize.insert(0, datasetroot.num_features)
NewNetworkSizeAdjust.append(NewNetworksize[0:-1])
print(NewNetworkSizeAdjust)
#OptimizedNet.apply(init_weights)
#OptimizedNet = DataParallel(OptimizedNet)
OptimizedNet = OptimizedNet.to(device)
cudnn.benchmark = True
criterionNew = nn.CrossEntropyLoss()
if optimizerName == "SGD":
optimizerNew = getattr(optim,
optimizerName)(OptimizedNet.parameters(),
lr=params[3],
momentum=0.9,
weight_decay=5e-4)
elif optimizerName == "Adam":
optimizerNew = getattr(optim,
optimizerName)(OptimizedNet.parameters(),
lr=params[3],
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=5e-4,
amsgrad=False)
TrainConvergence, TestConvergence, TestAcc = TrainPart(
modelName, datasetroot, start_epoch, num_epochs, trainloader,
testloader, OptimizedNet, optimizerNew, criterionNew, NumCutoff,
mark, True, model_to_save)
np.save(
"{}/{}Convergence/TrainConvergence-{}".format(
savepath, dataset, FileName), TrainConvergence)
np.save(
"{}/{}Convergence/TestConvergence-{}".format(
savepath, dataset, FileName), TestConvergence)
np.save(
"{}/{}Convergence/NewNetworkSizeAdjust-{}".format(
savepath, dataset, FileName), NewNetworkSizeAdjust)
#np.save(savepath+'TestConvergence-'+FileName,TestConvergence)
#torch.cuda.empty_cache()
print(
'dataset: {}, model name:{}, resized network size is {}, Number epoches:{}, Train error is: {}, Test error is: {}, test acc is {}\n'
.format(dataset, modelName, NewNetworksize[0:-1], num_epochs,
TrainConvergence[-1], TestConvergence[-1], TestAcc[-1]))
TestAccs.append(TestAcc)
np.save(
"{}/{}Convergence/MeanTestAccs-{}".format(savepath, dataset,
FileName), TestAccs)
print("The change of test error is:{}".format(TestAccs))
print_nvidia_useage()
num = 3
model_path = dir_path + "/cmodels/12_global_edge_attr_test/C_300.pt"
dataset_path = dir_path + '/dataset/cartesian/'
dataset_path += str(num) + "s/" if num != -1 else "all_nums"
def pf(data):
return data.y == num
pre_filter = pf if num != -1 else None
train_dataset = MNISTSuperpixels(dataset_path,
True,
pre_transform=T.Cartesian(),
pre_filter=pre_filter)
train_loader = DataLoader(train_dataset, batch_size=128)
print("loaded data")
pretrained_model = torch.load(model_path, map_location=device)
C = MoNet(25)
C.load_state_dict(pretrained_model.state_dict())
torch.save(pretrained_model.state_dict(),
"../mnist_superpixels/evaluation/C_state_dict.pt")
print("loaded model)")
# Testing Classification - remember to comment the first return in MoNet.forward()
def main(args):
args = init_dirs(args)
pt = T.Cartesian() if args.cartesian else T.Polar()
if args.dataset == 'sp':
train_dataset = MNISTSuperpixels(args.dataset_path, True, pre_transform=pt)
test_dataset = MNISTSuperpixels(args.dataset_path, False, pre_transform=pt)
train_loader = tgDataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
test_loader = tgDataLoader(test_dataset, batch_size=args.batch_size)
elif args.dataset == 'sm':
train_dataset = MNISTGraphDataset(args.dataset_path, args.num_hits, train=True)
train_loader = DataLoader(train_dataset, shuffle=True, batch_size=args.batch_size, pin_memory=True)
test_dataset = MNISTGraphDataset(args.dataset_path, args.num_hits, train=False)
test_loader = DataLoader(test_dataset, shuffle=True, batch_size=args.batch_size, pin_memory=True)
if(args.load_model):
C = torch.load(args.model_path + args.name + "/C_" + str(args.start_epoch) + ".pt").to(device)
else:
C = MoNet(args.kernel_size).to(device)
C_optimizer = torch.optim.Adam(C.parameters(), lr=args.lr, weight_decay=args.weight_decay)
if(args.scheduler):
C_scheduler = torch.optim.lr_scheduler.StepLR(C_optimizer, args.decay_step, gamma=args.lr_decay)
train_losses = []
test_losses = []
def plot_losses(epoch, train_losses, test_losses):
fig = plt.figure()
ax1 = fig.add_subplot(1, 2, 1)
ax1.plot(train_losses)
ax1.set_title('training')
ax2 = fig.add_subplot(1, 2, 2)
ax2.plot(test_losses)
ax2.set_title('testing')
plt.savefig(args.losses_path + args.name + "/" + str(epoch) + ".png")
plt.close()
def save_model(epoch):
torch.save(C, args.model_path + args.name + "/C_" + str(epoch) + ".pt")
def train_C(data, y):
C.train()
C_optimizer.zero_grad()
output = C(data)
# nll_loss takes class labels as target, so one-hot encoding is not needed
C_loss = F.nll_loss(output, y)
C_loss.backward()
C_optimizer.step()
return C_loss.item()
def test(epoch):
C.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data in test_loader:
if args.dataset == 'sp':
output = C(data.to(device))
y = data.y.to(device)
elif args.dataset == 'sm':
output = C(tg_transform(args, data[0].to(device)))
y = data[1].to(device)
test_loss += F.nll_loss(output, y, size_average=False).item()
pred = output.data.max(1, keepdim=True)[1]
correct += pred.eq(y.data.view_as(pred)).sum()
test_loss /= len(test_loader.dataset)
test_losses.append(test_loss)
print('test')
f = open(args.out_path + args.name + '.txt', 'a')
print(args.out_path + args.name + '.txt')
s = "After {} epochs, on test set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n".format(epoch, test_loss, correct, len(test_loader.dataset), 100. * correct / len(test_loader.dataset))
print(s)
f.write(s)
f.close()
for i in range(args.start_epoch, args.num_epochs):
print("Epoch %d %s" % ((i + 1), args.name))
C_loss = 0
test(i)
for batch_ndx, data in tqdm(enumerate(train_loader), total=len(train_loader)):
if args.dataset == 'sp':
C_loss += train_C(data.to(device), data.y.to(device))
elif args.dataset == 'sm':
C_loss += train_C(tg_transform(args, data[0].to(device)), data[1].to(device))
train_losses.append(C_loss / len(train_loader))
if(args.scheduler):
C_scheduler.step()
if((i + 1) % 10 == 0):
save_model(i + 1)
plot_losses(i + 1, train_losses, test_losses)
test(args.num_epochs)
import os.path as osp
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.datasets import MNISTSuperpixels
import torch_geometric.transforms as T
from torch_geometric.data import DataLoader
from torch_geometric.utils import normalized_cut
from torch_geometric.nn import (NNConv, graclus, max_pool, max_pool_x,
global_mean_pool)
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'MNIST')
train_dataset = MNISTSuperpixels(path, True, transform=T.Cartesian())
test_dataset = MNISTSuperpixels(path, False, transform=T.Cartesian())
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64)
d = train_dataset
def normalized_cut_2d(edge_index, pos):
row, col = edge_index
edge_attr = torch.norm(pos[row] - pos[col], p=2, dim=1)
return normalized_cut(edge_index, edge_attr, num_nodes=pos.size(0))
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
n1 = nn.Sequential(nn.Linear(2, 25), nn.ReLU(), nn.Linear(25, 32))
self.conv1 = NNConv(d.num_features, 32, n1)
def get_dataset(name,
sparse=True,
feat_str="deg+ak3+reall",
root=None,
aug=None,
aug_ratio=None):
if root is None or root == '':
path = osp.join(osp.expanduser('~'), 'pyG_data', name)
else:
path = osp.join(root, name)
degree = feat_str.find("deg") >= 0
onehot_maxdeg = re.findall("odeg(\d+)", feat_str)
onehot_maxdeg = int(onehot_maxdeg[0]) if onehot_maxdeg else None
k = re.findall("an{0,1}k(\d+)", feat_str)
k = int(k[0]) if k else 0
groupd = re.findall("groupd(\d+)", feat_str)
groupd = int(groupd[0]) if groupd else 0
remove_edges = re.findall("re(\w+)", feat_str)
remove_edges = remove_edges[0] if remove_edges else 'none'
edge_noises_add = re.findall("randa([\d\.]+)", feat_str)
edge_noises_add = float(edge_noises_add[0]) if edge_noises_add else 0
edge_noises_delete = re.findall("randd([\d\.]+)", feat_str)
edge_noises_delete = float(
edge_noises_delete[0]) if edge_noises_delete else 0
centrality = feat_str.find("cent") >= 0
coord = feat_str.find("coord") >= 0
pre_transform = FeatureExpander(degree=degree,
onehot_maxdeg=onehot_maxdeg,
AK=k,
centrality=centrality,
remove_edges=remove_edges,
edge_noises_add=edge_noises_add,
edge_noises_delete=edge_noises_delete,
group_degree=groupd).transform
print(aug, aug_ratio)
if 'MNIST' in name or 'CIFAR' in name:
if name == 'MNIST_SUPERPIXEL':
train_dataset = MNISTSuperpixels(path,
True,
pre_transform=pre_transform,
transform=T.Cartesian())
test_dataset = MNISTSuperpixels(path,
False,
pre_transform=pre_transform,
transform=T.Cartesian())
else:
train_dataset = ImageDataset(path,
name,
True,
pre_transform=pre_transform,
coord=coord,
processed_file_prefix="data_%s" %
feat_str)
test_dataset = ImageDataset(path,
name,
False,
pre_transform=pre_transform,
coord=coord,
processed_file_prefix="data_%s" %
feat_str)
dataset = (train_dataset, test_dataset)
else:
dataset = TUDatasetExt(path,
name,
pre_transform=pre_transform,
use_node_attr=True,
processed_filename="data_%s.pt" % feat_str,
aug=aug,
aug_ratio=aug_ratio)
dataset.data.edge_attr = None
return dataset
import os.path as osp
import torch
import torch.nn.functional as F
from torch_geometric.datasets import MNISTSuperpixels
from torch_geometric.data import DataListLoader
import torch_geometric.transforms as T
from torch_geometric.nn import SplineConv, global_mean_pool, DataParallel
path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data', 'MNIST')
dataset = MNISTSuperpixels(path, transform=T.Cartesian()).shuffle()
loader = DataListLoader(dataset, batch_size=1024, shuffle=True)
class Net(torch.nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = SplineConv(dataset.num_features, 32, dim=2, kernel_size=5)
self.conv2 = SplineConv(32, 64, dim=2, kernel_size=5)
self.lin1 = torch.nn.Linear(64, 128)
self.lin2 = torch.nn.Linear(128, dataset.num_classes)
def forward(self, data):
print('Inside Model: num graphs: {}, device: {}'.format(
data.num_graphs, data.batch.device))
x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr
x = F.elu(self.conv1(x, edge_index, edge_attr))
x = F.elu(self.conv2(x, edge_index, edge_attr))
x = global_mean_pool(x, data.batch)
x = F.elu(self.lin1(x))
start_time = time.time()
test_acc = validate(model, test_loader, device)
print('Test, Time: {:.4f}, Accuracy: {:.4f}'\
.format(time.time() - start_time, test_acc), file=file)
return model
if __name__ == '__main__':
# SuperPixels dataset
sp_path = os.path.join(os.path.dirname(os.path.realpath("/")),
"MNISTSuperpixel")
sp_test_dataset = MNISTSuperpixels(sp_path,
train=False,
transform=T.Cartesian())
sp_train_loader_25, sp_val_loader_25 = get_train_val_loader(
sp_path, train_batch_size=25, val_batch_size=25)
sp_train_loader_64, sp_val_loader_64 = get_train_val_loader(
sp_path, train_batch_size=64, val_batch_size=64)
sp_test_loader = DataLoader(sp_test_dataset, batch_size=1, shuffle=False)
# Skeletons dataset
sk_path = "dataset"
sk_train_dataset = MNISTSkeleton(sk_path, "train", transform=T.Polar())
sk_test_dataset = MNISTSkeleton(sk_path, "test", transform=T.Polar())
sk_val_dataset = MNISTSkeleton(sk_path, "val", transform=T.Polar())
sk_train_loader = DataLoader(sk_train_dataset, batch_size=64, shuffle=True)
sk_val_loader = DataLoader(sk_val_dataset, batch_size=64, shuffle=False)
parser.add_argument('--weight_decay', type=float, default=5e-4)
parser.add_argument('--batch_size', type=int, default=64)
parser.add_argument('--epochs', type=int, default=300)
parser.add_argument('--seed', type=int, default=1)
args = parser.parse_args()
args.data_fp = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data',
args.dataset)
device = torch.device('cuda', args.device_idx)
# deterministic
torch.manual_seed(args.seed)
cudnn.benchmark = False
cudnn.deterministic = True
train_dataset = MNISTSuperpixels(args.data_fp, True, pre_transform=T.Polar())
test_dataset = MNISTSuperpixels(args.data_fp, False, pre_transform=T.Polar())
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64)
def normalized_cut_2d(edge_index, pos):
row, col = edge_index
edge_attr = torch.norm(pos[row] - pos[col], p=2, dim=1)
return normalized_cut(edge_index, edge_attr, num_nodes=pos.size(0))
class MoNet(torch.nn.Module):
def __init__(self, kernel_size):
super(MoNet, self).__init__()
self.conv1 = GMMConv(1, 32, dim=2, kernel_size=kernel_size)
def main(args):
torch.manual_seed(4)
torch.autograd.set_detect_anomaly(True)
name = [args.name]
if WGAN:
name.append('wgan')
if GRU:
name.append('gru')
if GCNN:
name.append('gcnn')
# name.append('num_iters_{}'.format(args.num_iters))
# name.append('num_critic_{}'.format(args.num_critic))
args.name = '_'.join(name)
args.model_path = args.dir_path + '/models/'
args.losses_path = args.dir_path + '/losses/'
args.args_path = args.dir_path + '/args/'
args.figs_path = args.dir_path + '/figs/'
args.dataset_path = args.dir_path + '/dataset/'
args.err_path = args.dir_path + '/err/'
if(not exists(args.model_path)):
mkdir(args.model_path)
if(not exists(args.losses_path)):
mkdir(args.losses_path)
if(not exists(args.args_path)):
mkdir(args.args_path)
if(not exists(args.figs_path)):
mkdir(args.figs_path)
if(not exists(args.err_path)):
mkdir(args.err_path)
if(not exists(args.dataset_path)):
mkdir(args.dataset_path)
try:
# python2
file_tmp = urllib.urlretrieve(url, filename=None)[0]
except:
# python3
file_tmp = urllib.request.urlretrieve(url, filename=args.dataset)[0]
prev_models = [f[:-4] for f in listdir(args.args_path)] # removing .txt
if (args.name in prev_models):
print("name already used")
# if(not args.load_model):
# sys.exit()
else:
mkdir(args.losses_path + args.name)
mkdir(args.model_path + args.name)
mkdir(args.figs_path + args.name)
if(not args.load_model):
f = open(args.args_path + args.name + ".txt", "w+")
f.write(str(vars(args)))
f.close()
else:
f = open(args.args_path + args.name + ".txt", "r")
temp = args.start_epoch
args = eval(f.read())
f.close()
args.load_model = True
args.start_epoch = temp
# return args2
def pf(data):
return data.y == args.num
pre_filter = pf if args.num != -1 else None
print("loading")
# Change to True !!
X = SuperpixelsDataset(args.dataset_path, args.num_hits, train=TRAIN, num=NUM, device=device)
tgX = MNISTSuperpixels(args.dir_path, train=TRAIN, pre_transform=T.Cartesian(), pre_filter=pre_filter)
X_loaded = DataLoader(X, shuffle=True, batch_size=args.batch_size, pin_memory=True)
tgX_loaded = tgDataLoader(tgX, shuffle=True, batch_size=args.batch_size)
print("loaded")
if(args.load_model):
G = torch.load(args.model_path + args.name + "/G_" + str(args.start_epoch) + ".pt")
D = torch.load(args.model_path + args.name + "/D_" + str(args.start_epoch) + ".pt")
else:
G = Graph_Generator(args.node_feat_size, args.fe_hidden_size, args.fe_out_size, args.gru_hidden_size, args.gru_num_layers, args.num_iters, args.num_hits, args.dropout, args.leaky_relu_alpha, hidden_node_size=args.hidden_node_size, int_diffs=INT_DIFFS, gru=GRU, device=device).to(device)
if(GCNN):
D = MoNet(kernel_size=args.kernel_size, dropout=args.dropout, device=device).to(device)
# D = Gaussian_Discriminator(args.node_feat_size, args.fe_hidden_size, args.fe_out_size, args.gru_hidden_size, args.gru_num_layers, args.num_iters, args.num_hits, args.dropout, args.leaky_relu_alpha, kernel_size=args.kernel_size, hidden_node_size=args.hidden_node_size, int_diffs=INT_DIFFS, gru=GRU).to(device)
else:
D = Graph_Discriminator(args.node_feat_size, args.fe_hidden_size, args.fe_out_size, args.gru_hidden_size, args.gru_num_layers, args.num_iters, args.num_hits, args.dropout, args.leaky_relu_alpha, hidden_node_size=args.hidden_node_size, int_diffs=INT_DIFFS, gru=GRU, device=device).to(device)
print("Models loaded")
if(WGAN):
G_optimizer = optim.RMSprop(G.parameters(), lr=args.lr_gen)
D_optimizer = optim.RMSprop(D.parameters(), lr=args.lr_disc)
else:
G_optimizer = optim.Adam(G.parameters(), lr=args.lr_gen, weight_decay=5e-4)
D_optimizer = optim.Adam(D.parameters(), lr=args.lr_disc, weight_decay=5e-4)
print("optimizers loaded")
normal_dist = Normal(0, 0.2)
def wasserstein_loss(y_out, y_true):
return -torch.mean(y_out * y_true)
if(WGAN):
criterion = wasserstein_loss
else:
if(LSGAN):
criterion = torch.nn.MSELoss()
else:
criterion = torch.nn.BCELoss()
# print(criterion(torch.tensor([1.0]),torch.tensor([-1.0])))
def gen(num_samples, noise=0):
if(noise == 0):
noise = normal_dist.sample((num_samples, args.num_hits, args.hidden_node_size)).to(device)
return G(noise)
# transform my format to torch_geometric's
def tg_transform(X):
batch_size = X.size(0)
cutoff = 0.32178 # found empirically to match closest to Superpixels
pos = X[:, :, :2]
x1 = pos.repeat(1, 1, 75).reshape(batch_size, 75*75, 2)
x2 = pos.repeat(1, 75, 1)
diff_norms = torch.norm(x2 - x1 + 1e-12, dim=2)
diff = x2-x1
diff = diff[diff_norms < cutoff]
norms = diff_norms.reshape(batch_size, 75, 75)
neighborhood = torch.nonzero(norms < cutoff, as_tuple=False)
edge_attr = diff[neighborhood[:, 1] != neighborhood[:, 2]]
neighborhood = neighborhood[neighborhood[:, 1] != neighborhood[:, 2]] # remove self-loops
unique, counts = torch.unique(neighborhood[:, 0], return_counts=True)
edge_slices = torch.cat((torch.tensor([0]).to(device), counts.cumsum(0)))
edge_index = neighborhood[:, 1:].transpose(0, 1)
# normalizing edge attributes
edge_attr_list = list()
for i in range(batch_size):
start_index = edge_slices[i]
end_index = edge_slices[i+1]
temp = diff[start_index:end_index]
max = torch.max(temp)
temp = temp/(2*max + 1e-12) + 0.5
edge_attr_list.append(temp)
edge_attr = torch.cat(edge_attr_list)
x = X[:, :, 2].reshape(batch_size*75, 1)+0.5
pos = 27*pos.reshape(batch_size*75, 2)+13.5
zeros = torch.zeros(batch_size*75, dtype=int).to(device)
zeros[torch.arange(batch_size)*75] = 1
batch = torch.cumsum(zeros, 0)-1
return Batch(batch=batch, x=x, edge_index=edge_index.contiguous(), edge_attr=edge_attr, y=None, pos=pos)
def draw_graph(graph, node_r, im_px):
imd = im_px + node_r
img = np.zeros((imd, imd), dtype=np.float)
circles = []
for node in graph:
circles.append((draw.circle_perimeter(int(node[1]), int(node[0]), node_r), draw.disk((int(node[1]), int(node[0])), node_r), node[2]))
for circle in circles:
img[circle[1]] = circle[2]
return img
def save_sample_outputs(name, epoch, dlosses, glosses):
print("drawing figs")
fig = plt.figure(figsize=(10, 10))
num_ims = 100
node_r = 30
im_px = 1000
gen_out = gen(args.batch_size).cpu().detach().numpy()
for i in range(int(num_ims/args.batch_size)):
gen_out = np.concatenate((gen_out, gen(args.batch_size).cpu().detach().numpy()), 0)
gen_out = gen_out[:num_ims]
gen_out[gen_out > 0.47] = 0.47
gen_out[gen_out < -0.5] = -0.5
gen_out = gen_out*[im_px, im_px, 1] + [(im_px+node_r)/2, (im_px+node_r)/2, 0.55]
for i in range(1, num_ims+1):
fig.add_subplot(10, 10, i)
im_disp = draw_graph(gen_out[i-1], node_r, im_px)
plt.imshow(im_disp, cmap=cm.gray_r, interpolation='nearest')
plt.axis('off')
plt.savefig(args.figs_path + args.name + "/" + str(epoch) + ".png")
plt.close()
plt.figure()
plt.plot(dlosses, label='Discriminitive loss')
plt.plot(glosses, label='Generative loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.savefig(args.losses_path + args.name + "/" + str(epoch) + ".png")
plt.close()
print("saved figs")
def save_models(name, epoch):
torch.save(G, args.model_path + args.name + "/G_" + str(epoch) + ".pt")
torch.save(D, args.model_path + args.name + "/D_" + str(epoch) + ".pt")
# from https://github.com/EmilienDupont/wgan-gp
def gradient_penalty(real_data, generated_data):
batch_size = real_data.size()[0]
# Calculate interpolation
alpha = torch.rand(batch_size, 1, 1)
alpha = alpha.expand_as(real_data).to(device)
interpolated = alpha * real_data.data + (1 - alpha) * generated_data.data
interpolated = Variable(interpolated, requires_grad=True).to(device)
del alpha
torch.cuda.empty_cache()
# Calculate probability of interpolated examples
prob_interpolated = D(interpolated)
# Calculate gradients of probabilities with respect to examples
gradients = torch_grad(outputs=prob_interpolated, inputs=interpolated, grad_outputs=torch.ones(prob_interpolated.size()).to(device), create_graph=True, retain_graph=True, allow_unused=True)[0].to(device)
gradients = gradients.contiguous()
# Gradients have shape (batch_size, num_channels, img_width, img_height),
# so flatten to easily take norm per example in batch
gradients = gradients.view(batch_size, -1)
# Derivatives of the gradient close to 0 can cause problems because of
# the square root, so manually calculate norm and add epsilon
gradients_norm = torch.sqrt(torch.sum(gradients ** 2, dim=1) + 1e-12)
# Return gradient penalty
return args.gp_weight * ((gradients_norm - 1) ** 2).mean()
def train_D(data):
D.train()
D_optimizer.zero_grad()
run_batch_size = data.shape[0] if not GCNN else data.y.shape[0]
if(not WGAN):
Y_real = torch.ones(run_batch_size, 1).to(device)
Y_fake = torch.zeros(run_batch_size, 1).to(device)
try:
D_real_output = D(data)
gen_ims = gen(run_batch_size)
tg_gen_ims = tg_transform(gen_ims)
use_gen_ims = tg_gen_ims if GCNN else gen_ims
D_fake_output = D(use_gen_ims)
if(WGAN):
D_loss = D_fake_output.mean() - D_real_output.mean() + gradient_penalty(data, use_gen_ims)
else:
D_real_loss = criterion(D_real_output, Y_real)
D_fake_loss = criterion(D_fake_output, Y_fake)
D_loss = D_real_loss + D_fake_loss
D_loss.backward()
D_optimizer.step()
except:
print("Generated Images")
print(gen_ims)
print("Transformed Images")
print(tg_gen_ims)
print("Discriminator Output")
print(D_fake_output)
torch.save(gen_ims, args.err_path + args.name + "_gen_ims.pt")
torch.save(tg_gen_ims.x, args.err_path + args.name + "_x.pt")
torch.save(tg_gen_ims.pos, args.err_path + args.name + "_pos.pt")
torch.save(tg_gen_ims.edge_index, args.err_path + args.name + "_edge_index.pt")
return
return D_loss.item()
def train_G():
G.train()
G_optimizer.zero_grad()
if(not WGAN):
Y_real = torch.ones(args.batch_size, 1).to(device)
gen_ims = gen(args.batch_size)
if(GCNN):
gen_ims = tg_transform(gen_ims)
D_fake_output = D(gen_ims)
if(WGAN):
G_loss = -D_fake_output.mean()
else:
G_loss = criterion(D_fake_output, Y_real)
G_loss.backward()
G_optimizer.step()
return G_loss.item()
D_losses = []
G_losses = []
# save_models(name, 0)
# save_sample_outputs(args.name, 0, D_losses, G_losses)
# @profile
def train():
for i in range(args.start_epoch, args.num_epochs):
print("Epoch %d %s" % ((i+1), args.name))
D_loss = 0
G_loss = 0
loader = tgX_loaded if GCNN else X_loaded
for batch_ndx, data in tqdm(enumerate(loader), total=len(loader)):
if(batch_ndx > 0 and batch_ndx % (args.num_critic+1) == 0):
G_loss += train_G()
else:
D_loss += train_D(data.to(device)) if GCNN else train_D(data[0].to(device))
D_losses.append(D_loss/len(X_loaded)/2)
G_losses.append(G_loss/len(X_loaded))
if((i+1) % 5 == 0):
save_sample_outputs(args.name, i+1, D_losses, G_losses)
if((i+1) % 5 == 0):
save_models(args.name, i+1)
train()
