利用py文件训练

Kedreamix · Kedreamix · commit a6260f3916ea · 2022-05-08T16:16:24.000+08:00
diff --git a/CIFAR10_code/dataloader.py b/CIFAR10_code/dataloader.py
@@ -0,0 +1,32 @@
+import torch
+import torchvision
+import torchvision.transforms as transforms
+
+# Data
+def get_training_dataloader(batch_size = 64, num_workers = 4, shuffle = True):
+    print('==> Preparing Train data..')
+    transform_train = transforms.Compose([
+        transforms.RandomCrop(32, padding=4),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+    ])
+    trainset = torchvision.datasets.CIFAR10(
+        root='./data', train=True, download=True, transform=transform_train)
+    trainloader = torch.utils.data.DataLoader(
+        trainset, batch_size=batch_size, shuffle=shuffle, num_workers= num_workers)
+    return trainloader
+    
+def get_test_dataloader(batch_size = 64, num_workers = 4, shuffle = True): 
+    print('==> Preparing Test data..')   
+    transform_test = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+    ])
+
+    testset = torchvision.datasets.CIFAR10(
+        root='./data', train=False, download=True, transform=transform_test)
+    testloader = torch.utils.data.DataLoader(
+        testset, batch_size=batch_size, shuffle=shuffle, num_workers= num_workers)
+    return testloader
+
diff --git a/CIFAR10_code/nets/LeNet5.py b/CIFAR10_code/nets/LeNet5.py
@@ -0,0 +1,65 @@
+'''
+LetNet in Pytorch
+'''
+import torch
+import torch.nn as nn
+
+class LeNet5(nn.Module):
+    def __init__(self, num_classes = 10, init_weights=True):
+        super(LeNet5,self).__init__()
+        self.conv1 = nn.Sequential(
+            # 输入 32x32x3 -> 28x28x6 (32-5)/1 + 1=28
+            nn.Conv2d(in_channels=3,out_channels=6,kernel_size=5,stride=1),
+            nn.ReLU(),
+            # 经过2x2的maxpool，变成14x14 (28-2)/2+1
+            nn.MaxPool2d(kernel_size=2,stride=2)
+        )
+        
+        self.conv2 = nn.Sequential(
+            # 输入 14x14x6 ->  10x10x16    (14-5)/1 + 1 = 10
+            nn.Conv2d(in_channels=6,out_channels=16,kernel_size=5,stride=1),
+            nn.ReLU(),
+            # (10-2)/2 + 1 = 5 
+            nn.MaxPool2d(kernel_size=2,stride=2)
+        )
+        
+        self.fc = nn.Sequential(
+            nn.Linear(5*5*16,120),
+            nn.ReLU(),
+            nn.Linear(120,84),
+            nn.ReLU(),
+            nn.Linear(84,num_classes)
+        )
+        if init_weights:
+            self._initialize_weights()
+    def forward(self,x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        # 要把多维度的tensor展平成一维
+        x = x.view(x.size()[0],-1)
+        x = self.fc(x)
+        return x
+    
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                nn.init.constant_(m.bias, 0)    
+
+def test():
+    net = LeNet5()
+    x = torch.randn(2,3,32,32)
+    y = net(x)
+    print(y.size())
+    from torchinfo import summary
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    net = net.to(device)
+    summary(net,(2,3,32,32))
+    
diff --git a/CIFAR10_code/train.py b/CIFAR10_code/train.py
@@ -0,0 +1,205 @@
+'''Train CIFAR10 with PyTorch.'''
+import imp
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.backends.cudnn as cudnn
+import torchvision
+import torchvision.transforms as transforms
+
+import os
+import argparse
+from utils import get_acc,EarlyStopping
+from dataloader import get_test_dataloader, get_training_dataloader
+from tqdm import tqdm
+
+
+classes = ('plane', 'car', 'bird', 'cat', 'deer',
+           'dog', 'frog', 'horse', 'ship', 'truck')
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='PyTorch CIFAR10 Training')
+    parser.add_argument('--lr', default=0.1, type=float, help='learning rate')
+    parser.add_argument('--cuda', action='store_true', default=False, help =' use GPU?')
+    parser.add_argument('--batch-size', default=64, type=int, help = "Batch Size for Training")
+    parser.add_argument('--num-workers', default=2, type=int, help = 'num-workers')
+    parser.add_argument('--net', type = str, default='MobileNetv1', help='net type')
+    parser.add_argument('--epochs', type = int, default=20, help = 'Epochs')
+    parser.add_argument('--resume', '-r', action='store_true', help='resume from checkpoint')
+    parser.add_argument('--patience', '-p', type = int, default=7, help='patience for Early stop')
+    args = parser.parse_args()
+    
+    print(args)
+    best_acc = 0  # best test accuracy
+    start_epoch = 0  # start from epoch 0 or last checkpoint epoch
+    
+    # Train Data
+    trainloader = get_training_dataloader(batch_size = args.batch_size, num_workers = args.num_workers)
+    testloader = get_test_dataloader(batch_size = args.batch_size, num_workers = args.num_workers, shuffle=False)
+    # Model
+    print('==> Building model..')
+    if args.net == 'VGG16':
+        from nets.VGG import VGG
+        net = VGG('VGG16')
+    elif args.net == 'VGG19':
+        from nets.VGG import VGG
+        net = VGG('VGG19')
+    elif args.net == 'ResNet18':
+        from nets.ResNet import ResNet18
+        net = ResNet18()
+    elif args.net == 'ResNet34':
+        from nets.ResNet import ResNet34
+        net = ResNet34()
+    elif args.net == 'LeNet':
+        from nets.LeNet5 import LeNet5
+        net = LeNet5()
+    elif args.net == 'AlexNet':
+        from nets.AlexNet import AlexNet
+        net = AlexNet()
+    elif args.net == 'DenseNet':
+        from nets.DenseNet import densenet_cifar
+        net = densenet_cifar()
+    elif args.net == 'MobileNetv1':
+        from nets.MobileNetv1 import MobileNet
+        net = MobileNet()
+    elif args.net == 'MobileNetv2':
+        from nets.MobileNetv2 import MobileNetV2
+        net = MobileNetV2()
+
+    if args.cuda and torch.cuda.is_available():
+        device = 'cuda'
+        net = torch.nn.DataParallel(net)
+        cudnn.benchmark = True
+    else:
+        device = 'cpu'
+        
+    
+    if args.resume:
+        # Load checkpoint.
+        print('==> Resuming from checkpoint..')
+        assert os.path.isdir('checkpoint'), 'Error: no checkpoint directory found!'
+        checkpoint = torch.load('./checkpoint/{}_ckpt.pth'.format(args.net))
+        net.load_state_dict(checkpoint['net'])
+        best_acc = checkpoint['acc']
+        start_epoch = checkpoint['epoch']
+        args.lr = checkpoint['lr']
+
+    early_stopping = EarlyStopping(patience = args.patience, verbose=True)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(net.parameters(), lr=args.lr,
+                        momentum=0.9, weight_decay=5e-4)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)
+    
+    epochs = args.epochs
+    def train(epoch):
+        epoch_step = len(trainloader)
+        if epoch_step == 0:
+            raise ValueError("数据集过小，无法进行训练，请扩充数据集，或者减小batchsize")
+        net.train()
+        train_loss = 0
+        train_acc = 0
+        print('Start Train')
+        with tqdm(total=epoch_step,desc=f'Epoch {epoch + 1}/{epochs}',postfix=dict,mininterval=0.3) as pbar:
+            for step,(im,label) in enumerate(trainloader,start=0):
+                im = im.to(device)
+                label = label.to(device)
+                #---------------------
+                #  释放内存
+                #---------------------
+                if hasattr(torch.cuda, 'empty_cache'):
+                    torch.cuda.empty_cache()
+                #----------------------#
+                #   清零梯度
+                #----------------------#
+                optimizer.zero_grad()
+                #----------------------#
+                #   前向传播forward
+                #----------------------#
+                outputs = net(im)
+                #----------------------#
+                #   计算损失
+                #----------------------#
+                loss = criterion(outputs,label)
+                train_loss += loss.data
+                train_acc += get_acc(outputs,label)
+                #----------------------#
+                #   反向传播
+                #----------------------#
+                # backward
+                loss.backward()
+                # 更新参数
+                optimizer.step()
+                lr = optimizer.param_groups[0]['lr']
+                pbar.set_postfix(**{'Train Loss' : train_loss.item()/(step+1),
+                                    'Train Acc' :train_acc.item()/(step+1),  
+                                    'Lr'   : lr})
+                pbar.update(1)
+        # train_loss = train_loss.item() / len(trainloader)
+        # train_acc = train_acc.item() * 100 / len(trainloader)    
+        print('Finish Train')
+    def test(epoch):
+        global best_acc
+        epoch_step_test = len(testloader)
+        if epoch_step_test == 0:
+                raise ValueError("数据集过小，无法进行训练，请扩充数据集，或者减小batchsize")
+        
+        net.eval()
+        test_loss = 0
+        test_acc = 0
+        print('Start Test')
+        #--------------------------------
+        #   相同方法，同train
+        #--------------------------------
+        with tqdm(total=epoch_step_test,desc=f'Epoch {epoch + 1}/{epochs}',postfix=dict,mininterval=0.3) as pbar2:
+            for step,(im,label) in enumerate(testloader,start=0):
+                im = im.to(device)
+                label = label.to(device)
+                with torch.no_grad():
+                    if step >= epoch_step_test:
+                        break
+                    
+                    # 释放内存
+                    if hasattr(torch.cuda, 'empty_cache'):
+                        torch.cuda.empty_cache()
+                    #----------------------#
+                    #   前向传播
+                    #----------------------#
+                    outputs = net(im)
+                    loss = criterion(outputs,label)
+                    test_loss += loss.data
+                    test_acc += get_acc(outputs,label)
+                    
+                    pbar2.set_postfix(**{'Test Acc': test_acc.item()/(step+1),
+                                'Test Loss': test_loss.item() / (step + 1)})
+                    pbar2.update(1)
+        lr = optimizer.param_groups[0]['lr']
+        test_acc = test_acc.item() * 100 / len(testloader)
+        # Save checkpoint.
+        if test_acc > best_acc:
+            print('Saving..')
+            state = {
+                'net': net.state_dict(),
+                'acc': test_acc,
+                'epoch': epoch,
+                'lr': lr,
+            }
+            if not os.path.isdir('checkpoint'):
+                os.mkdir('checkpoint')
+            torch.save(state, './checkpoint/{}_ckpt.pth'.format(args.net))
+            best_acc = test_acc
+            
+        print('Finish Test')
+
+        early_stopping(test_loss, net)
+        # 若满足 early stopping 要求
+        if early_stopping.early_stop:
+            print("Early stopping")
+            # 结束模型训练
+            exit()
+        
+    for epoch in range(start_epoch, epochs):
+        train(epoch)
+        test(epoch)
+        scheduler.step()
+    torch.cuda.empty_cache()
+    
diff --git a/CIFAR10_code/utils.py b/CIFAR10_code/utils.py
