def train(dataset, dev):
trainNodes = torch.tensor(dataset.trainNodes, dtype=torch.long).to(dev)
trainNodesWithLabel = torch.tensor(dataset.trainNodesWithLabel,
dtype=torch.long).to(dev)
testNodes = torch.tensor(dataset.testNodes, dtype=torch.long).to(dev)
features = torch.tensor(dataset.features, dtype=torch.float32).to(dev)
labels = torch.tensor(dataset.labels, dtype=torch.float32).to(dev)
adj = None
if os.path.exists('adj_matrix/{}_adj.npy'.format(
dataset.datasetName)) is False:
adj = nx.adjacency_matrix(dataset.graph,
np.sort(list(dataset.graph.nodes))).A
adj = prerocess_adj(adj)
np.save('adj_matrix/{}_adj.npy'.format(dataset.datasetName), adj)
adj = torch.tensor(adj, dtype=torch.float32).to(dev)
else:
adj = torch.tensor(np.load('adj_matrix/{}_adj.npy'.format(
dataset.datasetName)),
dtype=torch.float32).to(dev)
Net = GCN(inDim=dataset.features.shape[1],
hidDim=16,
outDim=dataset.labels.shape[1],
numOfGCNLayers=2,
bias=False,
dropout=0.5)
Net.to(dev)
optimizer = torch.optim.Adam(Net.parameters(), lr=0.01,
weight_decay=5e-4) # adam optimizer
# ------------------------------------------ train ------------------------------------------ #
acc = 0.0
for epoch in range(200):
Net.train()
optimizer.zero_grad()
out = Net(features, adj)
loss = F.nll_loss(out[trainNodesWithLabel],
torch.argmax(labels[trainNodesWithLabel], dim=1))
loss.backward()
# print(epoch, loss.item())
optimizer.step()
with torch.no_grad():
Net.eval()
preds = Net(features, adj)
testPreds = torch.argmax(preds[testNodes], dim=1)
testLabels = torch.argmax(labels[testNodes], dim=1)
acc = (testPreds == testLabels).float().mean().item()
# print(epoch, acc)
print(acc)
return acc
def __init__(self, model_path: str, C: list, node_feature_dim: int ,num_class = 2, c=0, hyp1=1, hyp2=2, start=None, nfeat=7, dropout=0.1):
"""
:param C: Candidate set of nodes (list)
:param start: Starting node (defaults to randomised node)
"""
super(Generator, self).__init__()
self.nfeat = nfeat
self.dropout = dropout
self.c = c # c为要指定生成类别c的图
self.fc = nn.Linear(nfeat, 8)
self.gc1 = GraphConvolution(8, 16)
self.gc2 = GraphConvolution(16, 24)
self.gc3 = GraphConvolution(24, 32)
# MLP1
# 2 FC layers with hidden dimension 16
self.mlp1 = nn.Sequential(nn.Linear(32, 16), nn.Linear(16, 1))
# MLP2
# 2 FC layers with hidden dimension 24
self.mlp2 = nn.Sequential(nn.Linear(64, 24), nn.Linear(24, 1))
# Hyperparameters
self.hyp1 = hyp1
self.hyp2 = hyp2
self.candidate_set = C
# Default starting node (if any)
if start is not None:
self.start = start
self.random_start = False
else:
self.start = random.choice(np.arange(0, len(self.candidate_set)))
self.random_start = True
# Load GCN for calculating reward
self.model = GCN(nfeat=node_feature_dim,
nclass=num_class,
dropout=dropout)
self.model.load_state_dict(torch.load(model_path))
for param in self.model.parameters():
param.requires_grad = False
self.reset_graph()
def main():
# load dataset
features, labels, adj = load_cora()
# create graph convolutional network
gcn = GCN(input_dim=features.shape[-1],
hidden_dim=16,
output_dim=labels.shape[-1],
adj=adj)
gcnloss = Loss(gcn)
# the training set gives all labels, so it is a supervised learning.
mask = tf.ones((1, features.shape[-2]))
# train context
optimizer = tf.keras.optimizers.Adam(1e-2)
checkpoint = tf.train.Checkpoint(model=gcn,
optimizer=optimizer,
optimizer_step=optimizer.iterations)
checkpoint.restore(tf.train.latest_checkpoint('checkpoints'))
# create log
log = tf.summary.create_file_writer('checkpoints')
avg_loss = tf.keras.metrics.Mean(name='loss', dtype=tf.float32)
while True:
with tf.GradientTape() as tape:
outputs = gcn(features)
loss = gcnloss([outputs, labels, mask])
avg_loss.update_state(loss)
# apply gradients
grads = tape.gradient(loss, gcn.trainable_variables)
optimizer.apply_gradients(zip(grads, gcn.trainable_variables))
# write log
if tf.equal(optimizer.iterations % 10, 0):
with log.as_default():
tf.summary.scalar('loss',
avg_loss.result(),
step=optimizer.iterations)
print('Step #%d Loss: %.6f' %
(optimizer.iterations, avg_loss.result()))
avg_loss.reset_states()
# save checkpoint
checkpoint.save(join('checkpoints', 'ckpt'))
gcn.save_weights('gcn.h5')
def FedAvgGCNTrain(args):
os.environ['CUDA_VISIBLE_DEVICES'] = args['gpu']
# dev = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
dev = torch.device("cpu")
# dataset
dataset = dataGenerator(args['dataset'])
features = torch.tensor(dataset.features, dtype=torch.float32).to(dev)
labels = torch.tensor(dataset.labels, dtype=torch.float32).to(dev)
testNodes = torch.tensor(dataset.testNodes, dtype=torch.long).to(dev)
adj = None
if os.path.exists('adj_matrix/{}_adj.npy'.format(
dataset.datasetName)) is False:
adj = nx.adjacency_matrix(dataset.graph,
np.sort(list(dataset.graph.nodes))).A
adj = prerocess_adj(adj)
np.save('adj_matrix/{}_adj.npy'.format(dataset.datasetName), adj)
adj = torch.tensor(adj, dtype=torch.float32).to(dev)
else:
adj = torch.tensor(np.load('adj_matrix/{}_adj.npy'.format(
dataset.datasetName)),
dtype=torch.float32).to(dev)
# ------------------------------------- clients group -----------------------------------
clients = ClientsGroup(dataset, args['num_of_clients'], args['labelNIID'],
args['graphNIID'], dev)
numOfClientsPerComm = int(
max(args['num_of_clients'] * args['cFraction'], 1))
numOfClientsPerComm = numOfClientsPerComm - 1
withoutFedAvg = {'client0': 0, 'client1': 0, 'client2': 0, 'client3': 0}
withFedAvg = {'client0': 0, 'client1': 0, 'client2': 0, 'client3': 0}
newclient = 0.0
globalModel = 0.0
for _ in range(10):
# --------------------------------------- GCN Model --------------------------------------
Net = GCN(inDim=dataset.features.shape[1],
hidDim=args['hidden_layer_dim'],
outDim=dataset.labels.shape[1],
numOfGCNLayers=args['num_of_GCLayer'],
bias=args['bias'],
dropout=args['dropout'])
Net.to(dev)
lossFun = F.nll_loss # loss function
optimizer = torch.optim.Adam(Net.parameters(),
lr=args['learning_rate'],
weight_decay=5e-4) # adam optimizer
# initialParameter
initialParameter = {}
for key, var in Net.state_dict().items():
initialParameter[key] = var.clone()
print('============== without FedAvg 测试 ===============')
for id, client in clients.clientsSet.items():
localParameters, localTestSize = clients.clientsSet[
id].localUpdate(Net, lossFun, optimizer, initialParameter, 200)
localTestAcc = clients.clientsSet[id].localTest(
Net, localParameters)
withoutFedAvg[id] = withoutFedAvg[id] + localTestAcc
print('{} model'.format(id))
print('--local accuracy: {}%'.format(localTestAcc * 100))
# ------------------------------------ FedAvg GCN train ----------------------------------
print('============== with FedAvg 测试 ===============')
globalParameters = {}
for key, var in initialParameter.items():
globalParameters[key] = var.clone()
for i in range(args['num_comm']):
order = np.random.permutation(len(clients.clientsSet) - 1)
clientsIncomm = [
'client{}'.format(k) for k in order[0:numOfClientsPerComm]
]
sumParameters = None
totalSize = 0
for client in clientsIncomm:
localParameters, localTestSize = clients.clientsSet[
client].localUpdate(Net, lossFun, optimizer,
globalParameters, args['local_epoch'])
if sumParameters is None:
sumParameters = {}
for key, var in localParameters.items():
sumParameters[key] = var.clone() * localTestSize
else:
for var in sumParameters:
sumParameters[var] = sumParameters[
var] + localParameters[var] * localTestSize
totalSize = totalSize + localTestSize
for var in globalParameters:
globalParameters[var] = sumParameters[var] / totalSize
if (i + 1) == args['num_comm']:
for client in clientsIncomm:
localTestAcc = clients.clientsSet[client].localTest(
Net, globalParameters)
print('{} model'.format(client))
print('--local accuracy: {}%'.format(localTestAcc * 100))
withFedAvg[client] = withFedAvg[client] + localTestAcc
with torch.no_grad():
Net.load_state_dict(globalParameters, strict=True)
Net.eval()
preds = Net(features, adj)
testPreds = torch.argmax(preds[testNodes], dim=1)
testLabels = torch.argmax(labels[testNodes], dim=1)
acc = (testPreds == testLabels).float().mean().item()
if (i + 1) == args['num_comm']:
print('global model')
print("--global accuracy： {}%".format(acc * 100))
globalModel = globalModel + acc
print('================ New Client =================')
localTestAcc = clients.clientsSet['client{}'.format(
numOfClientsPerComm)].localTest(Net, globalParameters)
print('client{} model'.format(numOfClientsPerComm))
print("--local accuracy： {}%".format(localTestAcc * 100))
newclient = newclient + localTestAcc
for id, acc in withoutFedAvg.items():
print(id, acc / 10)
for id, acc in withFedAvg.items():
print(id, acc / 10)
print(globalModel / 10)
print(newclient / 10)
def FedAvgGCNTrain(args):
os.environ['CUDA_VISIBLE_DEVICES'] = args['gpu']
# dev = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
dev = torch.device("cpu")
# dataset
dataset = dataGenerator(args['dataset'])
features = torch.tensor(dataset.features, dtype=torch.float32).to(dev)
labels = torch.tensor(dataset.labels, dtype=torch.float32).to(dev)
testNodes = torch.tensor(dataset.testNodes, dtype=torch.long).to(dev)
adj = None
if os.path.exists('adj_matrix/{}_adj.npy'.format(
dataset.datasetName)) is False:
adj = nx.adjacency_matrix(dataset.graph,
np.sort(list(dataset.graph.nodes))).A
adj = prerocess_adj(adj)
np.save('adj_matrix/{}_adj.npy'.format(dataset.datasetName), adj)
adj = torch.tensor(adj, dtype=torch.float32).to(dev)
else:
adj = torch.tensor(np.load('adj_matrix/{}_adj.npy'.format(
dataset.datasetName)),
dtype=torch.float32).to(dev)
# ------------------------------------- clients group -----------------------------------
clients = ClientsGroup(dataset, args['num_of_clients'], args['labelNIID'],
args['graphNIID'], dev)
numOfClientsPerComm = int(
max(args['num_of_clients'] * args['cFraction'], 1))
numOfClientsPerComm = numOfClientsPerComm - 1 # remain the last one to be a new client
# --------------------------------------- GCN Model --------------------------------------
Net = GCN(inDim=dataset.features.shape[1],
hidDim=args['hidden_layer_dim'],
outDim=dataset.labels.shape[1],
numOfGCNLayers=args['num_of_GCLayer'],
bias=args['bias'],
dropout=args['dropout'])
Net.to(dev)
lossFun = F.nll_loss # loss function
optimizer = torch.optim.Adam(Net.parameters(),
lr=args['learning_rate'],
weight_decay=5e-4) # adam optimizer
# initialParameter
initialParameter = {}
for key, var in Net.state_dict().items():
initialParameter[key] = var.clone()
le = [1, 2, 5, 10, 20]
for e in le:
# ------------------------------------ FedAvg GCN train ----------------------------------
print('epoch {}'.format(e))
globalAcc = []
for id, client in clients.clientsSet.items():
client.testAcc = []
globalParameters = {}
for key, var in initialParameter.items():
globalParameters[key] = var.clone()
for i in range(args['num_comm']):
order = np.random.permutation(len(clients.clientsSet) - 1)
clientsIncomm = [
'client{}'.format(k) for k in order[0:numOfClientsPerComm]
]
sumParameters = None
totalSize = 0
for client in clientsIncomm:
localParameters, localTestSize = clients.clientsSet[
client].localUpdate(Net, lossFun, optimizer,
globalParameters, e)
if sumParameters is None:
sumParameters = {}
for key, var in localParameters.items():
sumParameters[key] = var.clone() * localTestSize
else:
for var in sumParameters:
sumParameters[var] = sumParameters[
var] + localParameters[var] * localTestSize
totalSize = totalSize + localTestSize
for var in globalParameters:
globalParameters[var] = sumParameters[var] / totalSize
for id, client in clients.clientsSet.items():
localTestAcc = client.localTest(Net, globalParameters)
with torch.no_grad():
Net.load_state_dict(globalParameters, strict=True)
Net.eval()
preds = Net(features, adj)
testPreds = torch.argmax(preds[testNodes], dim=1)
testLabels = torch.argmax(labels[testNodes], dim=1)
acc = (testPreds == testLabels).float().mean().item()
globalAcc.append(acc)
for id, client in clients.clientsSet.items():
path = 'result/IID/local_epoch/{}/{}'.format(
dataset.datasetName, id)
test_mkdir(path)
np.save(os.path.join(path, 'local_epoch_{}'.format(e)),
client.testAcc)
path = 'result/IID/local_epoch/{}/global'.format(dataset.datasetName)
test_mkdir(path)
np.save(os.path.join(path, 'local_epoch_{}'.format(e)), globalAcc)
class Generator(nn.Module):
def __init__(self, model_path: str, C: list, node_feature_dim: int ,num_class = 2, c=0, hyp1=1, hyp2=2, start=None, nfeat=7, dropout=0.1):
"""
:param C: Candidate set of nodes (list)
:param start: Starting node (defaults to randomised node)
"""
super(Generator, self).__init__()
self.nfeat = nfeat
self.dropout = dropout
self.c = c # c为要指定生成类别c的图
self.fc = nn.Linear(nfeat, 8)
self.gc1 = GraphConvolution(8, 16)
self.gc2 = GraphConvolution(16, 24)
self.gc3 = GraphConvolution(24, 32)
# MLP1
# 2 FC layers with hidden dimension 16
self.mlp1 = nn.Sequential(nn.Linear(32, 16), nn.Linear(16, 1))
# MLP2
# 2 FC layers with hidden dimension 24
self.mlp2 = nn.Sequential(nn.Linear(64, 24), nn.Linear(24, 1))
# Hyperparameters
self.hyp1 = hyp1
self.hyp2 = hyp2
self.candidate_set = C
# Default starting node (if any)
if start is not None:
self.start = start
self.random_start = False
else:
self.start = random.choice(np.arange(0, len(self.candidate_set)))
self.random_start = True
# Load GCN for calculating reward
self.model = GCN(nfeat=node_feature_dim,
nclass=num_class,
dropout=dropout)
self.model.load_state_dict(torch.load(model_path))
for param in self.model.parameters():
param.requires_grad = False
self.reset_graph()
def reset_graph(self):
"""
Reset g.G to default graph with only start node， 生成一个只有1个结点的图
"""
if self.random_start == True:
self.start = random.choice(np.arange(0, len(self.candidate_set)))
# 初始图除了第1个结点全被mask，这里由于邻接矩阵的边长为MAX_NUM_NODES + len(self.candidate_set)，所以mask的不仅为候选集结点，还有图中的所以虚结点
mask_start = torch.BoolTensor(
[False if i == 0 else True for i in range(MAX_NUM_NODES + len(self.candidate_set))])
adj = torch.zeros((MAX_NUM_NODES + len(self.candidate_set), MAX_NUM_NODES + len(self.candidate_set)),
dtype=torch.float32) # 这里adj shape为 [MAX_NUM_NODES + len(self.candidate_set), MAX_NUM_NODES + len(self.candidate_set)] 中间可能有空结点
feat = torch.zeros((MAX_NUM_NODES + len(self.candidate_set), len(self.candidate_set)), dtype=torch.float32)
feat[0, self.start] = 1
feat[np.arange(-len(self.candidate_set), 0), np.arange(0, len(self.candidate_set))] = 1
degrees = torch.zeros(MAX_NUM_NODES)
self.G = {'adj': adj, 'feat': feat, 'degrees': degrees, 'num_nodes': 1, 'mask_start': mask_start}
## 计算Gt->Gt+1
def forward(self, G_in):
## G_in为 Gt
G = copy.deepcopy(G_in)
x = G['feat'].detach().clone() # Gt的特征矩阵
adj = G['adj'].detach().clone() # Gt的邻接矩阵
## 对应 X = GCNs(Gt,C)
x = F.relu6(self.fc(x))
x = F.dropout(x, self.dropout, training=self.training)
x = F.relu6(self.gc1(x, adj))
x = F.dropout(x, self.dropout, training=self.training)
x = F.relu6(self.gc2(x, adj))
x = F.dropout(x, self.dropout, training=self.training)
x = F.relu6(self.gc3(x, adj))
x = F.dropout(x, self.dropout, training=self.training)
## pt,start=Softmax(MLPs(X))
p_start = self.mlp1(x)
p_start = p_start.masked_fill(G['mask_start'].unsqueeze(1), 0)
p_start = F.softmax(p_start, dim=0)
a_start_idx = torch.argmax(p_start.masked_fill(G['mask_start'].unsqueeze(1), -1))
## pt,end=Softmax(MLPs([X,x^start))
# broadcast
x1, x2 = torch.broadcast_tensors(x, x[a_start_idx])
x = torch.cat((x1, x2), 1) # cat increases dim from 32 to 64
# 计算maskt,end，除了候选集和Gt结点中未被选为初始结点的结点之外，其它均被mask
mask_end = torch.BoolTensor([True for i in range(MAX_NUM_NODES + len(self.candidate_set))])
mask_end[MAX_NUM_NODES:] = False
mask_end[:G['num_nodes']] = False
mask_end[a_start_idx] = True
p_end = self.mlp2(x)
p_end = p_end.masked_fill(mask_end.unsqueeze(1), 0)
p_end = F.softmax(p_end, dim=0)
a_end_idx = torch.argmax(p_end.masked_fill(mask_end.unsqueeze(1), -1))
# Return new G
# If a_end_idx is not masked, node exists in graph, no new node added
if G['mask_start'][a_end_idx] == False:
G['adj'][a_end_idx][a_start_idx] += 1
G['adj'][a_start_idx][a_end_idx] += 1
# Update degrees
G['degrees'][a_start_idx] += 1
G['degrees'][G['num_nodes']] += 1
else:
# Add node
G['feat'][G['num_nodes']] = G['feat'][a_end_idx]
# Add edge
G['adj'][G['num_nodes']][a_start_idx] += 1
G['adj'][a_start_idx][G['num_nodes']] += 1
# Update degrees
G['degrees'][a_start_idx] += 1
G['degrees'][G['num_nodes']] += 1
# Update start mask
G_mask_start_copy = G['mask_start'].detach().clone()
G_mask_start_copy[G['num_nodes']] = False
G['mask_start'] = G_mask_start_copy
G['num_nodes'] += 1
return p_start, a_start_idx, p_end, a_end_idx, G
### reward函数
def calculate_reward(self, G_t_1):
"""
Rtr Calculated from graph rules to encourage generated graphs to be valid
1. Only one edge to be added between any two nodes
2. Generated graph cannot contain more nodes than predefined maximum node number
3. (For chemical) Degree cannot exceed valency
If generated graph violates graph rule, Rtr = -1
Rtf Feedback from trained model
"""
rtr = self.check_graph_rules(G_t_1)
rtf = self.calculate_reward_feedback(G_t_1)
rtf_sum = 0
for m in range(rollout):
p_start, a_start, p_end, a_end, G_t_1 = self.forward(G_t_1)
rtf_sum += self.calculate_reward_feedback(G_t_1)
rtf = rtf + rtf_sum * self.hyp1 / rollout
return rtf + self.hyp2 * rtr
def calculate_reward_feedback(self, G_t_1):
"""
p(f(G_t_1) = c) - 1/l
where l denotes number of possible classes for f
"""
f = self.model(G_t_1['feat'], G_t_1['adj'])
return f[self.c] - 1 / len(f)
## graph rules
def check_graph_rules(self, G_t_1):
"""
For mutag, node degrees cannot exceed valency
"""
idx = 0
for d in G_t_1['degrees']:
if d is not 0:
node_id = torch.argmax(G_t_1['feat'][idx]) # Eg. [0, 1, 0, 0] -> 1
node = self.candidate_set[node_id] # Eg ['C.4', 'F.2', 'Br.7'][1] = 'F.2'
max_valency = int(node.split('.')[1]) # Eg. C.4 -> ['C', '4'] -> 4
# If any node degree exceeds its valency, return -1
if max_valency < d:
return -1
return 0
## 计算loss
def calculate_loss(self, Rt, p_start, a_start, p_end, a_end, G_t_1):
"""
Calculated from cross entropy loss (Lce) and reward function (Rt)
where loss = -Rt*(Lce_start + Lce_end)
"""
Lce_start = F.cross_entropy(torch.reshape(p_start, (1, 35)), a_start.unsqueeze(0))
Lce_end = F.cross_entropy(torch.reshape(p_end, (1, 35)), a_end.unsqueeze(0))
return -Rt * (Lce_start + Lce_end)
def load_model(self):
self.model = GCN(adj1=self.adjacency_mat1,
adj2=self.adjacency_mat2,
device=self.device).to(self.device)
self.loss = nn.CrossEntropyLoss().to(self.device)
class Processor():
def __init__(self):
self.get_argparser()
self.load_data()
self.adjacency_matrix()
self.load_model()
self.load_optimizer()
def get_argparser(self):
parser = argparse.ArgumentParser(description='Graph MNIST')
parser.add_argument(
'--batch-size',
type=int,
default=128,
help='input batch size for training (default: 128)')
parser.add_argument('--test-batch-size', type=int, default=1000)
parser.add_argument('--lr', type=int, default=0.01)
parser.add_argument('--momentum', type=int, default=0.5)
parser.add_argument('--epochs', type=int, default=10)
parser.add_argument('--log-interval', type=int, default=10)
parser.add_argument('--use-gpu',
action='store_true',
default=False,
help='use gpu True or False')
parser.add_argument('--gpu', type=int, default=3, help='GPU device')
#graph conv
parser.add_argument('--num_class', type=int, default=10)
parser.add_argument('--pool_size', default=(2, 2))
self.args = parser.parse_args()
if self.args.use_gpu:
if self.args.gpu == 0:
self.device = torch.device('cuda:0')
if self.args.gpu == 1:
self.device = torch.device('cuda:1')
if self.args.gpu == 2:
self.device = torch.device('cuda:2')
if self.args.gpu == 3:
self.device = torch.device('cuda:3')
else:
self.device = torch.device('cpu')
def load_data(self):
self.train_loader = torch.utils.data.DataLoader(
datasets.MNIST('./data',
train=True,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
])),
batch_size=self.args.batch_size,
shuffle=True,
num_workers=0)
self.test_loader = torch.utils.data.DataLoader(
datasets.MNIST('./data',
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
])),
batch_size=self.args.test_batch_size,
shuffle=False,
num_workers=0)
def load_model(self):
self.model = GCN(adj1=self.adjacency_mat1,
adj2=self.adjacency_mat2,
device=self.device).to(self.device)
self.loss = nn.CrossEntropyLoss().to(self.device)
def load_optimizer(self):
self.optimizer = optim.SGD(self.model.parameters(),
lr=self.args.lr,
momentum=self.args.momentum)
def train(self, epoch):
self.model.train()
for batch_idx, (data, label) in enumerate(
tqdm(self.train_loader, desc='batch', position=2)):
#get data
data = Variable(data.to(self.device))
label = Variable(label.to(self.device))
# forward
output = self.model(data)
loss = self.loss(output, label)
# backward
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
#print
if batch_idx % self.args.log_interval == 0:
self.print_log(
'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data),
len(self.train_loader.dataset),
100. * batch_idx / len(self.train_loader),
loss.item()),
print_tf=False)
def eval(self, epoch):
self.model.eval()
test_loss = 0
correct = 0
for batch_idx, (data, label) in enumerate(self.test_loader):
data = Variable(data.to(self.device))
label = Variable(label.to(self.device))
output = self.model(data)
test_loss += self.loss(output, label).item()
pred = output.max(
1, keepdim=True)[1] # get the index of the max log-probability
correct += pred.eq(label.view_as(pred)).sum().item()
test_loss /= len(self.test_loader.dataset)
#print
self.print_log('Test: Epoch[{}/{}], Accuracy: {}/{} ({:.2f}%)'.format(
epoch, self.args.epochs, correct, len(self.test_loader.dataset),
100. * correct / len(self.test_loader.dataset)),
print_tf=True)
def start(self):
for epoch in tqdm(range(1, self.args.epochs + 1),
desc='epoch',
position=1):
self.train(epoch)
self.eval(epoch)
def adjacency_matrix(self):
self.adjacency_mat1 = generate_adjacency_matrix(28)
self.adjacency_mat2 = generate_adjacency_matrix(14)
def print_log(self, str, print_tf):
#log
with open('{}/log.txt'.format('.'), 'a') as f:
f.write('\n' + str)
#print
if print_tf:
sys.stdout.write('\r{}'.format(str))
sys.stdout.flush()
else:
self.early_stop = True
else:
self.best_loss = loss
counter = 0
if __name__ == '__main__':
# adj_list: [188, 29, 29]
# features_list: [188, 29, 7]
# graph_labels: [188]
adj_list, features_list, graph_labels, idx_map, idx_train, idx_val, idx_test = load_split_MUTAG_data()
idx_train = torch.cat([idx_train, idx_val, idx_test])
model = GCN(nfeat=features_list[0].shape[1], # nfeat = 7
nclass=graph_labels.max().item() + 1, # nclass = 2
dropout=dropout)
optimizer = optim.Adam(model.parameters(),
lr=lr, weight_decay=weight_decay)
model.cuda()
features_list = features_list.cuda()
adj_list = adj_list.cuda()
graph_labels = graph_labels.cuda()
idx_train = idx_train.cuda()
idx_val = idx_val.cuda()
idx_test = idx_test.cuda()
# 训练模型
early_stopping = EarlyStopping(10, hit_min_before_stopping=True)
t_total = time.time()
from Load_dataset import load_split_MUTAG_data, accuracy
from Model import GCN
import time
import numpy as np
import torch
import torch.optim as optim
import torch.nn.functional as F
model_path = 'model/gcn_first.pth'
if __name__ == '__main__':
adj_list, features_list, graph_labels, idx_map, idx_train, idx_val, idx_test = load_split_MUTAG_data(
)
model = GCN(
nfeat=features_list[0].shape[1], # nfeat = 7
nclass=graph_labels.max().item() + 1, # nclass = 2
dropout=0.1)
model.eval()
outputs = []
for i in idx_test:
output = model(features_list[i], adj_list[i])
output = output.unsqueeze(0)
outputs.append(output)
output = torch.cat(outputs, dim=0)
loss_test = F.cross_entropy(output, graph_labels[idx_test])
acc_test = accuracy(output, graph_labels[idx_test])
print(loss_test)
print(acc_test)