pytorch--实现vgg

#!/usr/bin/env python3
# -*- coding:utf-8 -*- 
# @Time   : 2020/2/25 下午12:43
# @Author : MJ

import torch as t
import torch.nn as nn
import math
import torchvision.transforms as transforms
from torch.autograd import Variable
from torch.nn import functional as F
import torch.optim as optim
import torchvision as tv
# from torch.autograd import Variable

#define model
class VGG(nn.Module):
    def __init__(self,features,num_classes=10):
        super(VGG, self).__init__()
        # 网络结构（仅包含卷积层和池化层，不包含分类器）
        self.features = features

        self.classifer = nn.Sequential(
            #fc6
            nn.Linear(512,4096),
            nn.ReLU(),
            nn.Dropout(),

            #fc7
            nn.Linear(4096,4096),
            nn.ReLU(),
            nn.Dropout(),

            #fc8
            nn.Linear(4096,num_classes))

        #初始化权重
        self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0),-1)
        x =self.classifer(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m,nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0,math.sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m,nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m,nn.Linear):
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()

#生成网络每层的信息
def _make_layers(cfg,batch_norm=False):
    layers = []
    in_channels = 3
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2,stride=2)]
        else:
            #设定卷积层的输出数量
            conv2d = nn.Conv2d(in_channels,v,kernel_size=3,padding=1)
            if batch_norm:
                layers += [conv2d,nn.BatchNorm2d(v),nn.ReLU(inplace=True)]
            else:
                layers += [conv2d,nn.ReLU(inplace=True)]
            in_channels = v

    return nn.Sequential(*layers)

cfg = {
    'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}


def vgg11(**kwargs):
    model = VGG(_make_layers(cfg['A']),**kwargs)
    return model

def vgg11_bn(**kwargs):
    model = VGG(_make_layers(cfg['A'],batch_norm=True),**kwargs)
    return model

def vgg13(**kwargs):
    model = VGG(_make_layers(cfg['B']), **kwargs)
    return model


def vgg13_bn(**kwargs):
    model = VGG(_make_layers(cfg['B'], batch_norm=True), **kwargs)
    return model


def vgg16(**kwargs):
    model = VGG(_make_layers(cfg['D']), **kwargs)
    return model


def vgg16_bn(**kwargs):
    model = VGG(_make_layers(cfg['D'], batch_norm=True), **kwargs)
    return model


def vgg19(**kwargs):
    model = VGG(_make_layers(cfg['E']), **kwargs)
    return model


def vgg19_bn(**kwargs):
    model = VGG(_make_layers(cfg['E'], batch_norm=True), **kwargs)
    return model





#dataset function
def getData(BATCH_SIZE):
    transform_train = transforms.Compose([
        transforms.RandomResizedCrop(32),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485,0.459,0.467],std=[0.229,0.224,0.199]),])

    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485,0.495,0.512],std=[0.223,0.193,0.231]),])

    train_set = tv.datasets.CIFAR10(root='data/',train=True,transform=transform_train,download=True)
    train_loader = t.utils.data.DataLoader(train_set,batch_size=BATCH_SIZE,shuffle=True)

    test_set = tv.datasets.CIFAR10(root='data/',train=False,transform=transform_test,download=True)
    test_loader = t.utils.data.DataLoader(test_set,batch_size=BATCH_SIZE,shuffle=True)
    classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    return train_loader, test_loader, classes


#train function
def train(model,epoch,criterion,optimizer,data_loader):
    model.train()
    for batch_idx, (data,target) in enumerate(data_loader):
        model.cuda()
        data,target = Variable(data.cuda()),Variable(target.cuda())
        output = model(data)

        optimizer.zero_grad()
        loss = criterion(output,target)
        loss.backward()

        optimizer.step()

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

# test function
def test(model,epoch,criterion,data_loader):
    model.eval()
    test_loss = 0
    correct = 0
    for data, target in data_loader:
        model.cuda()
        data, target = Variable(data.cuda()),Variable(target.cuda())
        output = model(data)

        test_loss += criterion(output,target)
        pred = output.data.max(1)[1]
        correct += pred.eq(target.data).cpu().sum()

    test_loss /= len(data_loader)  # loss function already averages over batch size
    acc = correct / len(data_loader.dataset)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(data_loader.dataset), 100. * acc))

    return acc, test_loss


def defineModel():
    #define some hyperparameter
    BATCH_SIZE = 64
    LEARNING_RATE = 0.001
    NUM_EPOCHS = 5

    model = vgg16()
    model.cuda()

    criterion = nn.CrossEntropyLoss()
    criterion.cuda()

    optimizer = optim.Adam(model.parameters(),lr=LEARNING_RATE)

    return model, criterion, optimizer,NUM_EPOCHS, BATCH_SIZE

def mainfunction():

    model, criterion, optimizer, NUM_EPOCHS, BATCH_SIZE = defineModel()
    train_loader, test_loader, _ = getData(BATCH_SIZE)

    print("Training Model")
    print(model)
    for epoch in range(1, NUM_EPOCHS):
        train(model, epoch, criterion, optimizer, train_loader)
        with t.no_grad():
            acc, loss = test(model, 1, criterion, test_loader)
    t.save(model.state_dict(), 'checkpoints/VGG-001_net.t7')


if __name__ == '__main__':

    mainfunction()

参考：https://blog.youkuaiyun.com/zhanghao3389/article/details/85038252