limaosen0 / gxn Goto Github PK

for epoch in range(cmd_args.num_epochs):
        random.shuffle(train_idxes)
        classifier.train()
        avg_loss = loop_dataset(train_graphs, classifier, mi_loss, train_idxes, epoch, optimizer=optimizer, device=device)
        avg_loss[4] = 0.0
        print('\033[92maverage training of epoch %d: clsloss: %.5f miloss: %.5f loss %.5f acc %.5f auc %.5f\033[0m'
              % (epoch, avg_loss[0], avg_loss[1], avg_loss[2], avg_loss[3], avg_loss[4])) # noqa

        classifier.eval()
        test_loss = loop_dataset(test_graphs, classifier, mi_loss, list(range(len(test_graphs))), epoch, device=device)
        test_loss[4] = 0.0
        print('\033[93maverage test of epoch %d: clsloss: %.5f miloss: %.5f loss %.5f acc %.5f auc %.5f\033[0m'
              % (epoch, test_loss[0], test_loss[1], test_loss[2], test_loss[3], test_loss[4])) # noqa

        with open(logfile, 'a+') as log:
            log.write('test of epoch %d: clsloss: %.5f miloss: %.5f loss %.5f acc %.5f auc %.5f'
                      % (epoch, test_loss[0], test_loss[1], test_loss[2], test_loss[3], test_loss[4]) + '\n')

        if test_loss[3] > max_acc:
            max_acc = test_loss[3]
            fname = './checkpoint_%s/time_%s/FOLD%s/model_epoch%s.pt' % (cmd_args.data, first_timstr, foldidx, str(epoch))
            torch.save(classifier.state_dict(), fname)

with open('./result_%s/result_%s/acc_result_%s_%s.txt' % (cmd_args.data, first_timstr, cmd_args.data, first_timstr), 'a+') as f:
        f.write('\n')
        f.write('Fold index: ' + str(foldidx) + '\t')
        f.write(str(max_acc) + '\n')

class IndexSelect(nn.Module):

    def __init__(self, k, n_h, act,  R=1):
        super().__init__()
        self.k = k
        self.R = R
        self.sigm = nn.Sigmoid()
        self.fc = MLP(n_h, n_h, act)
        self.disc = Discriminator(n_h)
        self.gcn1 = GCN(n_h, n_h)

    def forward(self, seq1, seq2, A, samp_bias1=None, samp_bias2=None):
        h_1 = self.fc(seq1)
        h_2 = self.fc(seq2)
        h_n1 = self.gcn1(A, h_1)

        X = self.sigm(h_n1)
        ret, ret_true = self.disc(X, h_1, h_2, samp_bias1, samp_bias2)
        scores = self.sigm(ret_true).squeeze()
        num_nodes = A.shape[0]
        values, idx = torch.topk(scores, int(num_nodes))
        values1, idx1 = values[:int(self.k*num_nodes)], idx[:int(self.k*num_nodes)]
        values0, idx0 = values[int(self.k*num_nodes):], idx[int(self.k*num_nodes):]

        return ret, values1, idx1, idx0, h_n1