torchvizで計算グラフを可視化 - Qiita

import torch import torch.nn as nn import torch.nn.functional as F INPUT_SIZE = 10 class NeuralNet(nn.Module): def __init__(self): super(NeuralNet, self).__init__() self.fc1 = nn.Linear(INPUT, 256) self.fc2 = nn.Linear(256, 128) self.fc3 = nn.Linear(128, 10) def forward(self, x): x = self.fc1(x) x = F.relu(x) x = self.fc2(x) x = F.relu(x) x = self.fc3(x) x = F.softamx(x) return x torchvizによる計算グラフの