pytorch1.1.0-python3.6-CUDA9.0-以各种方式打印模型

CUDA模式下，当然也可以cpu模式下 

import tensorwatch as tw
dummy_input = torch.randn(1, 1, 32, 32, device='cuda')
tw_graph = tw.draw_model(sosnet32, dummy_input)
tw_graph.save("./mode_1.pdf")

import hiddenlayer as hl
from hiddenlayer import transforms as ht
dummy_input = torch.randn(1, 1, 32, 32, device='cuda')
hl_graph = hl.build_graph(sosnet32, dummy_input)
hl_graph.theme = hl.graph.THEMES["blue"].copy()
hl_graph.save("./model_2.pdf")

---

---

 

下面是我之前学习时的例子: 

from __future__ import print_function
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from efficientnet_pytorch import EfficientNet
from efficientnet_pytorch import utils

from torchsummary import summary
from torchstat import stat
from tensorboardX import SummaryWriter
writer = SummaryWriter('log')


from pytorchcv.model_provider import get_model as ptcv_get_model


def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        output1 = torch.nn.functional.log_softmax(output, dim=1)
        loss = F.nll_loss(output1, target)
        #loss = F.l1_loss(output, target)
        loss.backward()
        optimizer.step()

        #new ynh
        #每10个batch画个点用于loss曲线
        if batch_idx % 10 == 0:
            niter = epoch * len(train_loader) + batch_idx
            writer.add_scalar('Train/Loss', loss.data, niter)

        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))


def test(args, model, device, test_loader, epoch):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            output1 = torch.nn.functional.log_softmax(output, dim=1)
            test_loss += F.nll_loss(output1, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    # new ynh
    writer.add_scalar('Test/Accu', test_loss, epoch)


    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))


def main():
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--batch-size', type=int, default=10, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=10, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=10, metavar='N',
                        help='number of epochs to train (default: 10)')
    parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                        help='learning rate (default: 0.01)')
    parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                        help='SGD momentum (default: 0.5)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='how many batches to wait before logging training status')

    parser.add_argument('--save-model', action='store_true', default=False,
                        help='For Saving the current Model')
    args = parser.parse_args()
    use_cuda = not args.no_cuda and torch.cuda.is_available()

    torch.manual_seed(args.seed)

    device = torch.device("cuda" if use_cuda else "cpu")

    kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('./mnist', train=True, download=True,
                       transform=transforms.Compose([
                           transforms.Resize((224), interpolation=2),
                           transforms.Grayscale(3),
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=args.batch_size, shuffle=True, **kwargs)
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('./mnist', train=False, transform=transforms.Compose([
            transforms.Resize((224), interpolation=2),
            transforms.Grayscale(3),
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])),
        batch_size=args.test_batch_size, shuffle=True, **kwargs)

    blocks_args, global_params = utils.get_model_params('efficientnet-b0', override_params=None)
    #model = EfficientNet.from_pretrained('efficientnet-b0').to(device)#.cuda()
    #model = EfficientNet(blocks_args, global_params)#.to(device)  # .cuda()

    #dummy_input = torch.rand(1, 3, 224, 224)
    #writer.add_graph(model, (dummy_input,))

    #print(model)
    # cp resnet18-0982-0126861b.pth  ~/.torch/models
    model = ptcv_get_model("efficientnet_b1", pretrained=False).cuda()
    dummy_input = torch.randn(1, 3, 224, 224, device='cuda')
    torch.onnx.export(model, dummy_input, "./efficientnet_b0.onnx", verbose=True)

    #stat(model, (3, 224, 224))
    #model.to(device)
    #summary(model, (3, 224, 224))

    import hiddenlayer as hl
    from hiddenlayer import transforms as ht

    # no cuda()
    input = torch.zeros([1, 3, 224, 224], device='cuda')
    hl_graph = hl.build_graph(model, input)
    hl_graph.theme = hl.graph.THEMES["blue"].copy()
    #hl_graph.save("/home/boyun/PycharmProjects/EfficientNet-PyTorch/pytorch_resnet_bloks.pdf")
    dot = hl_graph.build_dot()
    dot.attr("graph", rankdir="TD")  # Topdown
    dot.format = "pdf"
    directory, file_name = "/home/boyun/PycharmProjects/EfficientNet-PyTorch/","efficientnet_b1.pdf"
    # Remove extension from file name. dot.render() adds it.
    file_name = file_name.replace("." + "pdf", "")
    dot.render(file_name, directory=directory, cleanup=True)

    print(model)
    print("-------------------------------------------")



    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        test(args, model, device, test_loader, epoch)

    if (args.save_model):
        torch.save(model.state_dict(), "mnist_cnn.pt")

    writer.close()


if __name__ == '__main__':
    main()

不同的版本，当时在一个一个试，仅供参考 

#-*-coding:utf-8-*-
from __future__ import print_function
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms, models
from efficientnet_pytorch import EfficientNet
from efficientnet_pytorch import utils
from MyDataset import MyDataset
from torchsummary import summary
from torchstat import stat
from tensorboardX import SummaryWriter
writer = SummaryWriter('log')

from torchviz import make_dot, make_dot_from_trace
import hiddenlayer as hl
from hiddenlayer import transforms as ht
import os
import torch.onnx

import tensorwatch as tw


def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    #for batch_idx, (data, target) in enumerate(train_loader):
    for batch_idx, data_ynh in enumerate(train_loader):
        # 获取图片和标签
        data, target = data_ynh
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        output1 = torch.nn.functional.log_softmax(output, dim=1)
        loss = F.nll_loss(output1, target)
        #loss = F.l1_loss(output, target)
        loss.backward()
        optimizer.step()

        #new ynh
        #每10个batch画个点用于loss曲线
        if batch_idx % 10 == 0:
            niter = epoch * len(train_loader) + batch_idx
            writer.add_scalar('Train/Loss', loss.data, niter)

        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))


def test(args, model, device, test_loader, epoch):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        #for data, target in test_loader:
        for data_ynh in test_loader:
            # 获取图片和标签
            data, target = data_ynh
            data, target = data.to(device), target.to(device)
            output = model(data)
            output1 = torch.nn.functional.log_softmax(output, dim=1)
            test_loss += F.nll_loss(output1, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    # new ynh
    writer.add_scalar('Test/Accu', test_loss, epoch)


    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))


def main():
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--batch-size', type=int, default=10, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=10, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=10, metavar='N',
                        help='number of epochs to train (default: 10)')
    parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                        help='learning rate (default: 0.01)')
    parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                        help='SGD momentum (default: 0.5)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='how many batches to wait before logging training status')

    parser.add_argument('--save-model', action='store_true', default=False,
                        help='For Saving the current Model')
    args = parser.parse_args()
    use_cuda = not args.no_cuda and torch.cuda.is_available()

    torch.manual_seed(args.seed)

    device = torch.device("cuda" if use_cuda else "cpu")

    kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}

    # -------------------------------------------- step 1/5 : 加载数据 -------------------------------------------
    train_txt_path = './Data/train.txt'
    valid_txt_path = './Data/valid.txt'
    # 数据预处理设置
    #normMean = [0.4948052, 0.48568845, 0.44682974]
    #normStd = [0.24580306, 0.24236229, 0.2603115]
    normMean = [104, 117, 123]
    normStd = [1, 1, 1]
    normTransform = transforms.Normalize(normMean, normStd)
    trainTransform = transforms.Compose([
        transforms.Resize(224),
        #transforms.RandomCrop(224, padding=4),
        transforms.ToTensor(),
        #normTransform
    ])

    validTransform = transforms.Compose([
        transforms.Resize(224),
        transforms.ToTensor(),
        #normTransform
    ])

    # 构建MyDataset实例 img_path是一种可在txt图片路径前面加入的一种机制
    train_data = MyDataset(img_path = '', txt_path=train_txt_path, transform=trainTransform)
    valid_data = MyDataset(img_path = '', txt_path=valid_txt_path, transform=validTransform)

    # 构建DataLoder
    train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size=16, shuffle=True, num_workers=2)
    valid_loader = torch.utils.data.DataLoader(dataset=valid_data, batch_size=16, num_workers=2)




    #blocks_args, global_params = utils.get_model_params('efficientnet-b0', override_params=None)
    #model = EfficientNet(blocks_args, global_params)
    model = EfficientNet.from_pretrained('efficientnet-b0').to(device)#.cuda()


    #dummy_input = torch.rand(1, 3, 224, 224).requires_grad_(True)
    #writer.add_graph(model, (dummy_input,))

    # no cuda()
    #vis_graph = make_dot(model(dummy_input), params=dict(model.named_parameters()))
    #vis_graph = make_dot(model(dummy_input), params=dict(list(model.named_parameters()) + [('x', dummy_input)]))
    #vis_graph.view()

    #no cuda()
    #input = torch.zeros([1, 3, 224, 224])
    #hl_graph = hl.build_graph(model, input)
    #hl_graph.theme = hl.graph.THEMES["blue"].copy()
    #hl_graph.save(os.path.join("/home/boyun/PycharmProjects/EfficientNet-1852", "pytorch_resnet_bloks.pdf"))



    #netron这个工具来可视化(读取ONNX文件)
    #https://discuss.pytorch.org/t/onnx-export-failed-couldnt-export-operator-aten-adaptive-avg-pool1d/30204
    #https://ptorch.com/news/95.html
    #model.train(False)
    #dummy_input = torch.randn(10, 3, 224, 224, device='cuda')
    #torch_out = torch.onnx._export(model, dummy_input, "./efficientnet-b0.onnx", export_params=True, verbose=True)


    #no .cuda()
    #tw_graph = tw.draw_model(model, [1, 3, 224, 224])

    #print(model)
    #model.to(device)
    #summary(model, (3, 224, 224))
    #model.cpu()
    #stat(model,  (3, 224, 224))
    print("-------------------------------------------")



    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        test(args, model, device, valid_loader, epoch)

    if (args.save_model):
        torch.save(model.state_dict(), "mnist_cnn.pt")

    writer.close()


if __name__ == '__main__':
    main()

或者：

from __future__ import print_function
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms, models
from efficientnet_pytorch import EfficientNet
from efficientnet_pytorch import utils

from torchsummary import summary
from torchstat import stat
from tensorboardX import SummaryWriter
writer = SummaryWriter('log')

import torch.onnx
import tensorwatch as tw

from torchviz import make_dot, make_dot_from_trace
import hiddenlayer as hl
from hiddenlayer import transforms as ht


def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        output1 = torch.nn.functional.log_softmax(output, dim=1)
        loss = F.nll_loss(output1, target)
        #loss = F.l1_loss(output, target)
        loss.backward()
        optimizer.step()

        #new ynh
        #每10个batch画个点用于loss曲线
        if batch_idx % 10 == 0:
            niter = epoch * len(train_loader) + batch_idx
            writer.add_scalar('Train/Loss', loss.data, niter)

        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))


def test(args, model, device, test_loader, epoch):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            output1 = torch.nn.functional.log_softmax(output, dim=1)
            test_loss += F.nll_loss(output1, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    # new ynh
    writer.add_scalar('Test/Accu', test_loss, epoch)


    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))


def main():
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--batch-size', type=int, default=10, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=10, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=10, metavar='N',
                        help='number of epochs to train (default: 10)')
    parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                        help='learning rate (default: 0.01)')
    parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                        help='SGD momentum (default: 0.5)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='how many batches to wait before logging training status')

    parser.add_argument('--save-model', action='store_true', default=False,
                        help='For Saving the current Model')
    args = parser.parse_args()
    use_cuda = not args.no_cuda and torch.cuda.is_available()

    torch.manual_seed(args.seed)

    device = torch.device("cuda" if use_cuda else "cpu")

    kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST(root='./mnist', train=True,download=True,
                       transform=transforms.Compose([
                           transforms.Resize((224), interpolation=2),
                           transforms.Grayscale(3),
                           transforms.ToTensor(),
                       ])),
        batch_size=args.batch_size, shuffle=True, **kwargs)
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST(root='./mnist', train=False, transform=transforms.Compose([
            transforms.Resize((224), interpolation=2),
            transforms.Grayscale(3),
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])),
        batch_size=args.test_batch_size, shuffle=True, **kwargs)

    blocks_args, global_params = utils.get_model_params('efficientnet-b0', override_params=None)
    #model = EfficientNet.from_pretrained('efficientnet-b0').to(device)#.cuda()
    model = EfficientNet(blocks_args, global_params)#.to(device)  # .cuda()

    #dummy_input = torch.rand(1, 3, 224, 224)
    #writer.add_graph(model, (dummy_input,))

    #dummy_input = torch.randn(10, 3, 224, 224, device='cuda')
    #model = model.cuda()
    model1 = models.alexnet(pretrained=True)#.cuda()
    #torch.onnx.export(model1, dummy_input, "efficientnet.onnx", verbose=True)

    #print(model)
    #x = x.view(x.size(0), 256 * 6 * 6)
    #x = x.view(x.size(0), -1)  这句话的出现就是为了将前面多维度的tensor展平成一维
    #x = x.view(x.size(0), -1)简化x = x.view(batchsize, -1)
    #AdaptiveAvgPool2d - [-1, 256, 6, 6]
    #Dropout - [-1, 9216]

    tw_graph = tw.draw_model(model1, [1, 3, 224, 224])
    tw_graph.save("/home/boyun/PycharmProjects/EfficientNet-py1.0.1/pytorch_resnet_bloks.pdf")

    input = torch.zeros([1, 3, 224, 224])
    hl_graph = hl.build_graph(model1, input)
    hl_graph.theme = hl.graph.THEMES["blue"].copy()
    hl_graph.save("/home/boyun/PycharmProjects/EfficientNet-py1.0.1/pytorch_resnet_bloks_1.pdf")

    #stat(model, (3, 224, 224))
    model.to(device)
    #summary(model, (3, 224, 224))

    print("-------------------------------------------")



    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        test(args, model, device, test_loader, epoch)

    if (args.save_model):
        torch.save(model.state_dict(), "mnist_cnn.pt")

    writer.close()


if __name__ == '__main__':
    main()