龙良曲pytorch学习笔记_ResNet18

main----dataloader----train----test

相对LeNet5的主函数来讲，仅仅是更换了模型名称，其他部分没有变化。

import torch
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision import transforms
from torch import nn,optim
from resnet import ResNet18

def main():
    batch_size = 32
    cifar_train = datasets.CIFAR10('cifar',train = True,transform = transforms.Compose([
        transforms.Resize((32,32)),
        transforms.ToTensor()
    ]),download = True)
    
    # 可以同时加载多张图片
    cifar_train = DataLoader(cifar_train,batch_size = batch_size,shuffle = True)
    
    cifar_test = datasets.CIFAR10('cifar',train = False,transform = transforms.Compose([
        transforms.Resize((32,32)),
        transforms.ToTensor(),
        # transforms.RandomRotation(5),
        transforms.Normalize(mean = [0.485,0,456,0,406],
                             std = [0.229.0.224.0.225])
    ]),download = True)
    
    # 可以同时加载多张图片
    cifar_test = DataLoader(cifar_test,batch_size = batch_size,shuffle = True)

    # 数据加载成功后可以检验shape
    x,label = iter(cifar_train).next()
    print('x:',x.shape,'label:',label.shape)

    device = torch.device('cuda')
    model = ResNet18().to(device)
    criteon = nn.CrossEntropyLoss().to(device)
    optimizer = optim.Adam(model.parameters(),lr=1e-3)
    
    print(model)
    
    for epoch in range(1000):
        
        model.train()
        for batchidx,(x,label) in enumerate(cifar_train):
            # x: [b,3,32,32], label: [b]
            x,label = x.to(device),label.to(device)
            
            logits = model(x)
            # logits:[b,10]
            # label:[b]
            loss = criteon(logits,label)
            
            # backprop
            optimizer.zero_grad()
            loss.backwark()
            optimizer.step()
            
        #
        print(epoch,loss.item())
        
        model.eval()
        # 不需要做梯度相关计算
        with torch.nn_grad():
            # test
            total_correct = 0
            total_num = 0
            for x,label in cifar_test:
                x,label = x.to(device),label.to(device)
                # logits:[b,10]
                logits = model(x)
                pred = logits.argmax(dim=1)
                # 获取一个batch的在累加
                total_correct = += torch.eq(pred,label).float().sum().item()
                # x.size(0)就是batch_size
                total_num += x.size(0)
                
            acc = total_correct / total_num
            print(epoch,acc)
            
if __name__ == '__main__'
    main()

ResNet18

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

class ResBlk(nn.Module):

    def __init__(self,ch_in,ch_out,stride = 1):
        super(ResBlk,self).__init__()
        
        # 改变stride是为了使得图片的size变小，以避免占用过多内存
        self.conv1 = nn.Conv2d(ch_in,ch_out,kernel_size = 3,stride = stride,padding = 1)
        self.bn1 = nn.BatchNorm2d(ch_out)
        self.conv2 = nn.Conv2d(ch_out,ch_out,kernel_size = 3,stride = 1,padding = 1)
        self.bn2 = nn.BatchNorm2d(ch_out)
        
        self.extra = nn.Squential()
        if ch_out != ch_in:
            # [b,ch_in,h,w] -->  [b,ch_in,h,w]
            self.extra = nn.Squential(
                # x要和f(x)的size也一样，所以也要设置stride
                # 而channel通过一个卷积层来使得他们一致
                nn.Conv2d(ch_in,ch_out,kernel_size = 1,stride = stride)
                nn.BatchNorm2d(ch_out)
            )
        
    def forward(self,x):
        out = F.relu(self.bn1(self.conv1(x)))
        # 这里的relu取决于自己
        out = F.relu(self.bn2(self.conv2(out)))
        # short cut
        # extra module: [b,ch_in,h,w] -->  [b,ch_in,h,w]
        # element-wise add 需要ch_in和ch_out相等
        # 由于是残差网络，所以要把f(x)和短路的x相加
        out = self.extra(x) + out
        
        return out
        
class ResNet18(nn.Module):
    
    def __init__(self):
        super(ResNet18,self).__init__()
        
        self.conv1 = nn.Sequential(
            nn.Conv2d(3,64,kernel_size = 3,stride = 1,padding = 1),
            nn.BatchNorm2d(64)
        )
        # followws 4 blocks
        # [b,64,h,w] --> [b,128,h,w]
        self.blk1 = ResBlk(64,128,stride = 2)
        # [b,128,h,w] --> [b,256,h,w]
        self.blk2 = ResBlk(128,256,stride = 2)
        # [b,256,h,w] --> [b,512,h,w]
        self.blk3 = ResBlk(256,512,stride = 2)
        # [b,512,h,w] --> [b,512,h,w]
        self.blk4 = ResBlk(512,512,stride = 2)
        
        # 线性层的输入需要测试之后才能知道
        self.outlayer = nn.Linear(512*1*1,10)
        
    def forward(self,x):
        x = F.relu(self.conv1(x))
        # [b,64,h,w] --> [b,1024,h,w]
        x = self.blk1(x)
        x = self.blk2(x)
        x = self.blk3(x)
        x = self.blk4(x)
        
        # print('after conv:',x.shape) # [b,512,2,2]
        # [b,512,1,1] --> [b,512,1,1]
        x = F.adaptive_avg_pool2d(x,[1,1])
        # print('after conv:',x.shape)
        x = x.view(x.size(0),-1)
        x = self.outlayer(x)
        
        return x
        
def main():
    
    blk = ResBlk(64,128,stride = 4)
    tmp = torch.randn(2,64,32,32)
    out = blk(tmp)
    print('block:',out.shape)
    
    x = torch.randn(2,3,32,32)
    model = ResNet18()
    out = model(x)
    print('resnet:',out.shape)