一个 torchsummary实例

torchsummary实例

这次需要使用torchsummary的时候突然忘记了向量格式，记录一下以免以后还忘记

import torch
import os
from torch import nn
from torch.nn import functional as F
from torch.autograd import Variable
import matplotlib.pyplot as plt
from torchvision.datasets import ImageFolder
import torch.optim as optim
import torch.utils.data
from PIL import Image
import torchvision.transforms as transforms
import copy
from torchsummary import summary


class MZSRModel(nn.Module):
    def __init__(self, channels, num_of_layers=8):
        super(MZSRModel, self).__init__()
        kernel_size = 3                         #卷积核大小为3
        padding = 1                             #填充1    
        features = 64                           #特征图维度=64
        layers = []
        layers.append(nn.Conv2d(in_channels=channels, out_channels=features, kernel_size=kernel_size, padding=padding, bias=False))
        layers.append(nn.ReLU(inplace=True))                    #3x3 same卷积+Relu激活
        for _ in range(num_of_layers-2):                        #16层，每层都是3x3卷积+BN+RELU
            layers.append(nn.Conv2d(in_channels=features, out_channels=features, kernel_size=kernel_size, padding=padding, bias=False))
            #layers.append(nn.BatchNorm2d(features))
            layers.append(nn.ReLU(inplace=True))
        layers.append(nn.Conv2d(in_channels=features, out_channels=channels, kernel_size=kernel_size, padding=padding, bias=False))
        self.features = nn.Sequential(*layers)
    def forward(self, x):
        x2=copy.deepcopy( x)
        out = self.features(x)+x2
        return out

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # PyTorch v0.4.0
model = MZSRModel(3).to(device)
summary(model,(3,10,10))