ts

zysubtle · zysubtle · commit 665e8fedd6f9 · 2022-12-25T21:17:05.000+08:00
diff --git a/res50_backbone.py b/res50_backbone.py
@@ -0,0 +1,106 @@
+import torch.nn as nn
+import torch
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, in_channel, out_channel, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel,
+                               kernel_size=1, stride=1, bias=False)  # squeeze channels
+        self.bn1 = nn.BatchNorm2d(out_channel)
+        # -----------------------------------------
+        self.conv2 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel,
+                               kernel_size=3, stride=stride, bias=False, padding=1)
+        self.bn2 = nn.BatchNorm2d(out_channel)
+        # -----------------------------------------
+        self.conv3 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel*self.expansion,
+                               kernel_size=1, stride=1, bias=False)  # unsqueeze channels
+        self.bn3 = nn.BatchNorm2d(out_channel*self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+
+    def forward(self, x):
+        identity = x
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, blocks_num, num_classes=1000, include_top=True):
+        super(ResNet, self).__init__()
+        self.include_top = include_top
+        self.in_channel = 64
+
+        self.conv1 = nn.Conv2d(3, self.in_channel, kernel_size=7, stride=2,
+                               padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(self.in_channel)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, blocks_num[0])
+        self.layer2 = self._make_layer(block, 128, blocks_num[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, blocks_num[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, blocks_num[3], stride=2)
+        if self.include_top:
+            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))  # output size = (1, 1)
+            self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+
+    def _make_layer(self, block, channel, block_num, stride=1):
+        downsample = None
+        if stride != 1 or self.in_channel != channel * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.in_channel, channel * block.expansion, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(channel * block.expansion))
+
+        layers = []
+        layers.append(block(self.in_channel, channel, downsample=downsample, stride=stride))
+        self.in_channel = channel * block.expansion
+
+        for _ in range(1, block_num):
+            layers.append(block(self.in_channel, channel))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        if self.include_top:
+            x = self.avgpool(x)
+            x = torch.flatten(x, 1)
+            x = self.fc(x)
+
+        return x
+
+
+def resnet50(num_classes=1000, include_top=True):
+    return ResNet(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, include_top=include_top)
diff --git a/ssd_model.py b/ssd_model.py
@@ -0,0 +1,46 @@
+import torch
+from torch import nn, Tensor
+
+
+
+class Backbone(nn.Module):
+    def __init__(self, pretrain_path=None):
+        super(Backbone, self).__init__()  # Q：子类调用父类的方法
+        net = resnet50()
+        self.out_channels = [1024, 512, 512, 256, 256, 256]  # 6个特征图的channel数目
+
+        # 使用预训练模型的权重
+        if pretrain_path is not None:
+            net.load_state_dict(torch.load(pretrain_path))
+        
+        # 从resnet中取前7个模块作为feature extractor
+        # 参考：https://blog.csdn.net/pengchengliu/article/details/113878358
+        self.feature_extractor = nn.Sequential(*list(net.children())[:7])
+
+        # 对Conv4的第1个Block进行修改
+        conv4_block1 = self.feature_extractor[-1][0]
+        conv4_block1.conv1.stride = (1, 1)
+        conv4_block1.conv2.stride = (1, 1)
+        conv4_block1.downsample[0].stride = (1, 1)
+
+    def forward(self, x):
+        x = self.feature_extractor(x)
+        return x
+
+
+class SSD300(nn.Module):
+    def __init__(self, backbone=None, num_classes=21):
+        super(SSD300, self).__init__()
+        if backbone==None:
+            raise Exception('backbone is None')
+        if not hasattr(backbone, 'out_channels'):
+            raise Exception('backbone not has attribute: out_channels')
+        self.feature_extractor = backbone
+
+        self.num_classes = num_classes
+        # out_channels = [1024, 512, 512, 256, 256, 256]
+        self._build_additional_features(self.feature_extractor.out_channels)
+    
+
+    def _build_additional_features(self, input_size):
+        
diff --git a/test.py b/test.py
@@ -0,0 +1,10 @@
+import torch
+from torch import nn, Tensor
+from torch.jit.annotations import List
+
+from res50_backbone import resnet50
+
+net = resnet50()
+
+
+print(*list(net.children())[:7])
diff --git a/train_ssd300.py b/train_ssd300.py
@@ -1,6 +1,160 @@
+import os
+import datetime
+
+import torch
+
+import transforms
+from my_dataset import VOCDataSet
+
+
+def create_model(num_classes=21):
+    backbone = Backbone()  # 特征提取器
+    model = SSD300(backbone=backbone, num_classes=num_classes)
+
+    pre_ssd_path = '***'
+    if os.path.exists(pre_ssd_path) is False:
+        raise FileNotFoundError('*** not found in {}'.format(pre_ssd_path))
+    pre_model_dict = torch.load(pre_ssd_path, map_location='cpu')
+    pre_weights_dict = pre_model_dict['model']
+
+    # 不加载类别预测器的权重，因为是voc和coco不同，但是可以使用回归预测器的权重
+    del_conf_loc_dict = {}
+    for k, v in pre_weights_dict.items():
+        split_key = k.split('.')
+        if 'conf' in k:
+            continue
+        del_conf_loc_dict.update({k: v})
+    
+    missing_keys, unexpected_keys = model.load_state_dict(del_conf_loc_dict, strict=False)
+    if len(missing_keys) != 0 or len(unexpected_keys) != 0:
+        print("missing_keys: ", missing_keys)
+        print("unexpected_keys: ", unexpected_keys)
+
+    return model    
+
 
 def main(parser_data):
+    # 指定GPU或CPU，Q：如何进行多GPU训练
+    device = torch.device(parser_data.device if torch.cuda.is_available() else 'cpu')
+    print('Using device {} training.'.format(device.type))
+
+    if not os.path.exists('save_weights'):
+        os.makedirs('save_weights')
+    
+    results_file = 'results{}.txt'.format(datetime.datetime.now().strftime('%Y%m%d-%H%M%S'))
+
+    data_transform = {
+        'train': transforms.Compose([
+            transforms.SSDCropping(),  # Q: 裁剪应该旨在增加训练样本，裁剪后GT也要改变？
+            transforms.Resize(),
+            transforms.ColorJitter(),
+            transforms.ToTensor(),
+            transforms.RandomHorizontalFlips(),
+            transforms.Normalization(),
+            # ???：输出default box集合的正样本和负样本(default_box的位置固定，正样本和GT匹配最佳且满足IOU>0.5)
+            transforms.AssignGTtoDefaultBox()
+        ]),
+        'val': transforms.Compose([
+            transforms.Resize(),
+            transforms.ToTensor(),
+            transforms.Normalization()
+        ])
+    }
+
+    VOC_root = parser_data.data_path
+    if os.path.exists(os.path.join(VOC_root, 'VOCdevkit')) is False:
+        raise FileNotFoundError('VOCdevkit does not in path: {}'.format(VOC_root))
+    
+    # VOCdevkit/VOC2012/ImageSets/Main/train.txt
+    train_dataset = VOCDataset(VOC_root, '2012', data_transform['train'], train_set='train.txt')
+    # 训练时batch size必须大于1.
+    batch_size = parser_data.batch_size
+    assert batch_size > 1, 'batch size must be greater than 1'  # Q：assert条件不满足，是否就推出程序？
+    # 防止最后一个batch_size=1，如果是就舍去
+    drop_last = True if len(train_dataset) % batch_size == 1 else False
+    nw = min([os.cput_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
+    print('Using %g dataloader workers' % nw)
+    train_data_loader = torch.utils.data.DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=nw,
+        collate_fn=train_dataset.collate_fn,  # 核对函数
+        drop_last=drop_last
+    )
+
+    # VOCdevkit/VOC2012/ImageSets/Main/val.txt
+    val_dataset = VOCDataset(VOC_root, '2007', data_transform['val'], train_set='val.txt')
+    val_data_loader = torch.utils.data.Dataloader(
+        val_dataset,
+        batch_size=batch_size,
+        num_workers=nw,
+        collate_fn=train_dataset.collate_fn  # ?：为何不是val_dataset.collate_fn
+    )
+
+    model = create_model(num_classes=parser_data.num_classes+1)
+    model.to(device)
+
+    # define optimizer
+    params = [p for p in model.parameters() if p.requires_grad]  # Q：可以更新的参数部分？
+    optimizer = torch.optim.SGD(params, lr=0.0005, momentum=0.9, weight_decay=0.0005)
+
+    # learning date scheduler，Q：step_size指每个epoch增大5倍？
+    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.3)
+
+    # 如果指定了上次训练保存的权重文件地址，则续接上次结果训练
+    if parser_data.resume != '':
+        check_point = torch.load(parser_data.resume, map_location='cpu')  # Q：在cpu上加载保存于GPU的模型？
+        model.load_state_dict(check_point['model'])
+        optimizer.load_state_dict(check_point['optimizer'])
+        lr_scheduler.load_state_dict(check_point['lr_scheduler'])
+        parser_data.start_epoch = check_point['epoch'] + 1
+        print('the training process from epoch {}...'.format(parser_data.start_epoch))
+    
+    train_loss = []
+    learning_rate = []
+    val_map = []
+
+    # 提前加载验证集数据，以免每次验证都重新加载一次
+    val_data = get_coco_api_from_dataset(val_data_loader.dataset)  # Q：这里是说一次性加载所有的数据？
+    for epoch in range(parser_data.start_epoch, parser_data.epochs):
+        mean_loss, lr = utils.train_one_epoch(model=model,optimizer=optimizer, data_loader=train_data_loader, device=device, epoch=epoch, print_freq=50)
+        train_loss.append(mean_loss.item())
+        learning_rate.append(lr)
+
+        # update learning rate
+        lr_scheduler.step()
+
+        # Q：是否使用data_set参数后就是一直把数据保存在GPU上，从而忽略了data_loader这种以batch加载数据的方式
+        coco_info = utils.evaluate(model=model, data_loader=val_data_loader, device=device, data_set=val_data)
+    
+        # write info txt
+        with open(results_file, 'a') as f:
+            result_info = [str(round(i, 4)) for i in coco_info + [mean_loss.item()]] + [str(round(lr, 6))]
+            txt = "epoch:{} {}".format(epoch, '  '.join(result_info))
+            f.write(txt + "\n")
+    
+        val_map.append(coco_info[1])
+
+        # save weights
+        save_files = {
+            'model': model.state_dict(),  # Q：state_dict()是什么意思？
+            'optimizer': optimizer.state_dict(),
+            'lr_scheduler': lr_scheduler.state_dict(),
+            'epoch': epoch
+        }
+
+        torch.save(save_files, './save_weights/ssd300-{}.pth'.format(epoch))
+    
+    # plot loss and lr curve
+    if len(train_loss) != 0 and len(learning_rate) != 0:
+        from plot_curve import plot_loss_and_lr
+        plot_loss_and_lr(train_loss, lr)
     
+    # plot mAP curve
+    if len(val_map) != 0:
+        from plot_curve import plot_map
+        plot_map(val_map)
 
 
 if __name__ == '__main__':
