GCNN実装にあたってのグラフとラベル

Question

##実装したい内容
dglライブラリを用いて、
グラフ構造の畳み込みニューラルネットワークを実装したいです。
２種類のグラフ(circle,star)
をランダムに生成して、２クラス分類をするものを作成する予定です。

モデルの定義と訓練は以下のようになります。
（pytorchを用いて実装しています。）

import dgl
import torch

def collate(samples):
    # The input `samples` is a list of pairs
    #  (graph, label).
    graphs, labels = map(list, zip(*samples))
    batched_graph = dgl.batch(graphs)
    return batched_graph, torch.tensor(labels)

from dgl.nn.pytorch import GraphConv

import torch.nn as nn
import torch.nn.functional as F

#モデルの定義
class Classifier(nn.Module):
    def __init__(self, in_dim, hidden_dim, n_classes):
        super(Classifier, self).__init__()
        self.conv1 = GraphConv(in_dim, hidden_dim)
        self.conv2 = GraphConv(hidden_dim, hidden_dim)
        self.classify = nn.Linear(hidden_dim, n_classes)
        #self.conv1,self.conv2をconvolution層として表現
        #nn.Linearで全結合層として表現

        
    #グラフの畳み込みと活性化関数
    #順伝播
    def forward(self, g):
        # Use node degree as the initial node feature. For undirected graphs, the in-degree
        # is the same as the out_degree.
        h = g.in_degrees().view(-1, 1).float()
        # Perform graph convolution and activation function.
        h = F.relu(self.conv1(g, h))
        h = F.relu(self.conv2(g, h))
        g.ndata['h'] = h
        # Calculate graph representation by averaging all the node representations.
        hg = dgl.mean_nodes(g, 'h')
        return self.classify(hg)

import torch.optim as optim
from torch.utils.data import DataLoader

#320個のトレーニングセットと、80個の訓練セットを作成する。
trainset = testclusternotest.MiniGCDatasetnotest(320, 10, 20)
testset = testclusternotest.MiniGCDatasetnotest(80, 10, 20)
# Use PyTorch's DataLoader and the collate function
# defined before.

data_loader = DataLoader(trainset, batch_size=32, shuffle=True,
                         collate_fn=collate)


# modelを作成
model = Classifier(1, 256, trainset.num_classes)
loss_func = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

model.train()

epoch_losses = []

#80回同じ操作を繰り返す。
for epoch in range(80):
    epoch_loss = 0
    for iter, (bg, label) in enumerate(data_loader):
        prediction = model(bg)
        loss = loss_func(prediction, label)
        #勾配のリセット
        optimizer.zero_grad()
        #逆伝播の更新
        loss.backward()
        #重みの更新
        optimizer.step()
        epoch_loss += loss.detach().item()
    # epoch_loss = epoch_loss/(iter+1)
    epoch_loss /= (iter + 1)
    print('Epoch {}, loss {:.4f}'.format(epoch, epoch_loss))
    epoch_losses.append(epoch_loss)

#使用したデータセット
以下のようなMiniGCDatasetと呼ばれるDGLライブラリ付属のデータセットを改変いたしました。元は８クラス分類でしたが、２クラスに書き換えました。

元のデータセットは、以下のURLからご覧ください。
https://docs.dgl.ai/en/0.4.x/_modules/dgl/data/minigc.html#MiniGCDataset

"""作成したデータセット"""
import math, os
import networkx as nx
import numpy as np

from .dgl_dataset import DGLDataset
from .utils import save_graphs, load_graphs, makedirs
from .. import backend as F
from ..convert import from_networkx
from ..transform import add_self_loop

__all__ = ['MiniGCDatasetnotest']

class MiniGCDatasetnotest(DGLDataset):

    def __init__(self, num_graphs, min_num_v, max_num_v, seed=0,
                 save_graph=True, force_reload=False, verbose=False):
        self.num_graphs = num_graphs
        self.min_num_v = min_num_v
        self.max_num_v = max_num_v
        self.seed = seed
        self.save_graph = save_graph
        super(MiniGCDatasetnotest, self).__init__(name="testclusternotest", hash_key=(num_graphs, min_num_v, max_num_v, seed),
                                            force_reload=force_reload,
                                            verbose=verbose)

    def process(self):
        self.graphs = []
        self.labels = []
        self._generate(self.seed)

    def __len__(self):
        """Return the number of graphs in the dataset."""
        return len(self.graphs)

    def __getitem__(self, idx):
        """Get the idx-th sample.

        Parameters
        ---------
        idx : int
            The sample index.

        Returns
        -------
        (:class:`dgl.Graph`, Tensor)
            The graph and its label.
        """
        return self.graphs[idx], self.labels[idx]

    def has_cache(self):
        graph_path = os.path.join(self.save_path, 'dgl_graph_{}.bin'.format(self.hash))
        if os.path.exists(graph_path):
            return True

        return False

    def save(self):
        """save the graph list and the labels"""
        if self.save_graph:
            graph_path = os.path.join(self.save_path, 'dgl_graph_{}.bin'.format(self.hash))
            save_graphs(str(graph_path), self.graphs, {'labels': self.labels})

    def load(self):
        graphs, label_dict = load_graphs(os.path.join(self.save_path, 'dgl_graph_{}.bin'.format(self.hash)))
        self.graphs = graphs
        self.labels = label_dict['labels']

    @property
    def num_classes(self):
        """Number of classes."""
        return 2

    def _generate(self, seed):
        if seed is not None:
            np.random.seed(seed)
        self._gen_cycle(self.num_graphs // 2)
        self._gen_star(self.num_graphs // 2)

        # preprocess
        for i in range(self.num_graphs):
            # convert to DGLGraph, and add self loops
            self.graphs[i] = add_self_loop(from_networkx(self.graphs[i]))
        self.labels = F.tensor(np.array(self.labels).astype(np.int))

    def _gen_cycle(self, n):
        for _ in range(n):
            num_v = np.random.randint(self.min_num_v, self.max_num_v)
            g = nx.cycle_graph(num_v)
            self.graphs.append(g)
            self.labels.append(0)

    def _gen_star(self, n):
        for _ in range(n):
            num_v = np.random.randint(self.min_num_v, self.max_num_v)
            # nx.star_graph(N) gives a star graph with N+1 nodes
            g = nx.star_graph(num_v - 1)
            self.graphs.append(g)
            self.labels.append(1)

##発生したエラー
Index Error: Target 5 is out of bounds.が出力されました。

IndexError                                Traceback (most recent call last)
<ipython-input-15-dfdf585f0b2f> in <module>
     36     for iter, (bg, label) in enumerate(data_loader):
     37         prediction = model(bg)
---> 38         loss = loss_func(prediction, label)
     39         #勾配のリセット
     40         optimizer.zero_grad()

~/anaconda/envs/pytorch/lib/python3.6/site-packages/torch/nn/functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
   2216                          .format(input.size(0), target.size(0)))
   2217     if dim == 2:
-> 2218         ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   2219     elif dim == 4:
   2220         ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)

IndexError: Target 5 is out of bounds.

#自分で試したこと
print関数にて、labelがきちんと出力されているのか以下で確認しました。

for epoch in range(80):
    epoch_loss = 0
    for iter, (bg, label) in enumerate(data_loader):
        prediction = model(bg)
        print(label)
        loss = loss_func(prediction, label)

出力結果

tensor([5, 1, 4, 3, 7, 4, 0, 6, 0, 0, 4, 0, 2, 3, 0, 3, 4, 5, 3, 1, 3, 4, 7, 7, 4, 6, 5, 7, 3, 0, 1, 1])

これより、tensorの最初の要素5が、0か1でないために、エラーが出ていると推測されます。元は８クラス分類だったのでこの出力は納得がいくのですが、どのようにデータセットを書き換えれば良いかがわかりません。
どなたかご教授願えませんでしょうか。