SE-Resnext网络搭建及预训练模型

SE-Resnext预训练

完整代码

# -*- coding: utf-8 -*-

import os, torch, glob, time, copy
import numpy as np
from torch.autograd import Variable
from PIL import Image  
from torchvision import models, transforms
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
import torch.optim as optim
from torch.optim import lr_scheduler
import pretrainedmodels
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
train_data_dir = '/home/train/'
test_data_dir = '/home/test/'
 
def get_lists(data_dir):        
    files_list = []
    labels_list = []
    sub_dirs = [x[0] for x in os.walk(data_dir) ]
    sub_dirs = sub_dirs[1:]
    
    extentions = ['jpg', 'jpeg', 'JPG', 'JPEG']
        
    for sub_dir in sub_dirs:
        if os.path.basename(sub_dir) == '0':
            label = [0]
        elif os.path.basename(sub_dir) == '1':
            label = [1] 
        elif os.path.basename(sub_dir) == '2':
            label = [2]
        elif os.path.basename(sub_dir) == '3':
            label = [3] 
        elif os.path.basename(sub_dir) == '4':
            label = [4]
        elif os.path.basename(sub_dir) == '5':
            label = [5] 
        elif os.path.basename(sub_dir) == '6':
            label = [6]
        elif os.path.basename(sub_dir) == '7':
            label = [7] 
        elif os.path.basename(sub_dir) == '8':
            label = [8]
        elif os.path.basename(sub_dir) == '9':
            label = [9] 
        elif os.path.basename(sub_dir) == '10':
            label = [10]
        elif os.path.basename(sub_dir) == '11':
            label = [11] 
        elif os.path.basename(sub_dir) == '12':
            label = [12]
        elif os.path.basename(sub_dir) == '13':
            label = [13] 
        elif os.path.basename(sub_dir) == '14':
            label = [14]
        elif os.path.basename(sub_dir) == '15':
            label = [15] 
        elif os.path.basename(sub_dir) == '16':
            label = [16]
        elif os.path.basename(sub_dir) == '17':
            label = [17]            
        for extention in extentions:
            file_glob = os.path.join(sub_dir, '*.' + extention)
            new_files = glob.glob(file_glob)
            files_list.extend(new_files)
            labels_list = labels_list + label * len(new_files)
    return files_list, labels_list
 
class myDataset(Dataset):
    # TensorDataset继承Dataset, 重载了__init__, __getitem__, __len__
    # 实现将一组Tensor数据对封装成Tensor数据集
    # 能够通过index得到数据集的数据，能够通过len，得到数据集大小
 
    def __init__(self, files_list, labels_list, transform):
        self.files_list = files_list
        self.labels_list = labels_list
        self.transform = transform
 
    def __getitem__(self, index):
        img = Image.open(self.files_list[index])
        img = self.transform(img)
        return img, self.labels_list[index]
 
    def __len__(self):
        return len(self.labels_list)
def make_save_dir(dirname,epoch):
    save_dir = os.path.join(dirname,epoch)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    return save_dir
def get_filename_for_saving(save_dir):
    return os.path.join(save_dir,"torh.pth")
def train_model(model, criterion, optimizer, scheduler, num_epochs=15):
    since1 = time.time()
 
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
 
    for epoch in range(num_epochs):
        since2 = time.time()
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)
 
        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                scheduler.step()
                model.train(True)  # Set model to training mode
            else:
                model.train(False)  # Set model to evaluate mode
 
            running_loss = 0.0
            running_corrects = 0
 
            # Iterate over data.
            for data in (dataloaders[phase]):
                # get the inputs
                inputs, labels = data
 
                # wrap them in Variable
                if use_gpu:
                    inputs = Variable(inputs.cuda())
                    labels = Variable(labels.cuda())
                else:
                    inputs, labels = Variable(inputs), Variable(labels)
 
                # zero the parameter gradients
                optimizer.zero_grad()
                # forward
                outputs = model(inputs)
                _, preds = torch.max(outputs.data, 1)
                loss = criterion(outputs, labels)
 
                # backward + optimize only if in training phase
                if phase == 'train':
                    loss.backward()
                    optimizer.step()
 
                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))
            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
                
        print(time.time() - since2)
        save_dir="/home/"+str(epoch)+".pth"
        torch.save(model,save_dir)
 
 
    time_elapsed = time.time() - since1
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))
 
    # load best model weights
    model.load_state_dict(best_model_wts)
    torch.save(model.state_dict(),'/home/11.13_resnet50_no.pth')
    return model
 
if __name__ == '__main__':
    BATCH_SIZE =32
    use_gpu = torch.cuda.is_available()

    transform = transforms.Compose([
            transforms.Scale(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485,0.456,0.406],
                                 std=[0.229,0.224,0.225])]
    )
    train_images, train_labels = get_lists(train_data_dir)
    test_images, test_labels = get_lists(test_data_dir)

    train_dataset = myDataset(train_images, train_labels, transform)
    test_dataset = myDataset(test_images, test_labels, transform)
    train_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE,
                               shuffle=True, num_workers=0)
    test_loader = DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE,
                               shuffle=False, num_workers=0)
    dataloaders = {}
    dataloaders['train'] = train_loader
    dataloaders['val'] = test_loader

    dataset_sizes = {}
    dataset_sizes['train'] = len(train_labels)
    dataset_sizes['val'] = len(test_labels)

    # model_ft = models.resnet50(pretrained=True)
    # # for param in model_ft.parameters():
    # #     param.requires_grad = False
    # num_ftrs = model_ft.fc.in_features
    # model_ft.fc = nn.Linear(num_ftrs, 18)

    model_name = 'se_resnext101_32x4d' # could be fbresnet152 or inceptionresnetv2
    model_ft = pretrainedmodels.__dict__[model_name](num_classes=1000, pretrained='imagenet')
    # num_ftrs = model_ft.fc.in_features
    model_ft.fc = nn.Linear(2048, 18)
    model_ft.eval()
    if use_gpu:
        model_ft = model_ft.cuda()
        model_ft = torch.nn.DataParallel(model_ft, device_ids=[0,1])

    criterion = nn.CrossEntropyLoss()

    # Observe that all parameters are being optimized
    # optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
    optimizer_ft = optim.Adam(model_ft.parameters(), lr=0.001, betas=(0.9,0.999),eps=1e-08,weight_decay=0,amsgrad=False)
    # Decay LR by a factor of 0.1 every 7 epochs
    exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
    model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,
                           num_epochs=25)
#    for (batch_x, batch_y) in enumerate(train_loader):
#        print(batch_x, batch_y)

1、预训练模型下载
点击这里下载
2、pytorch每训练一轮保存一次模型

def train_model(model, criterion, optimizer, scheduler, num_epochs=15):
    since1 = time.time()
 
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
 
    for epoch in range(num_epochs):
        since2 = time.time()
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)
 
        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                scheduler.step()
                model.train(True)  # Set model to training mode
            else:
                model.train(False)  # Set model to evaluate mode
 
            running_loss = 0.0
            running_corrects = 0
 
            # Iterate over data.
            for data in (dataloaders[phase]):
                # get the inputs
                inputs, labels = data
 
                # wrap them in Variable
                if use_gpu:
                    inputs = Variable(inputs.cuda())
                    labels = Variable(labels.cuda())
                else:
                    inputs, labels = Variable(inputs), Variable(labels)
 
                # zero the parameter gradients
                optimizer.zero_grad()
                # forward
                outputs = model(inputs)
                _, preds = torch.max(outputs.data, 1)
                loss = criterion(outputs, labels)
 
                # backward + optimize only if in training phase
                if phase == 'train':
                    loss.backward()
                    optimizer.step()
 
                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))
            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
                
        print(time.time() - since2)
        save_dir="/home/vivo3/11.14se_resnext101_32x4d/"+str(epoch)+".pth"//每训练一轮保存一次模型
        torch.save(model,save_dir)
 
 
    time_elapsed = time.time() - since1
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))
 
    # load best model weights
    model.load_state_dict(best_model_wts)
    torch.save(model.state_dict(),'/home/11.13_resnet50_no.pth')//保存最佳模型
    return model
 

3、

SE-Resnext网络搭建完整代码

'''
New for ResNeXt:
1. Wider bottleneck
2. Add group for conv2
'''

import torch.nn as nn
import math

__all__ = ['SE_ResNeXt', 'se_resnext_50', 'se_resnext_101', 'se_resnext_152']


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None, num_group=32):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes*2, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes*2)
        self.conv2 = nn.Conv2d(planes*2, planes*2, kernel_size=3, stride=stride,
                               padding=1, bias=False, groups=num_group)
        self.bn2 = nn.BatchNorm2d(planes*2)
        self.conv3 = nn.Conv2d(planes*2, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

        if planes == 64:
            self.globalAvgPool = nn.AvgPool2d(56, stride=1)
        elif planes == 128:
            self.globalAvgPool = nn.AvgPool2d(28, stride=1)
        elif planes == 256:
            self.globalAvgPool = nn.AvgPool2d(14, stride=1)
        elif planes == 512:
            self.globalAvgPool = nn.AvgPool2d(7, stride=1)
        self.fc1 = nn.Linear(in_features=planes * 4, out_features=round(planes / 4))
        self.fc2 = nn.Linear(in_features=round(planes / 4), out_features=planes * 4)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        residual = 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)

        if self.downsample is not None:
            residual = self.downsample(x)

        original_out = out
        out = self.globalAvgPool(out)
        out = out.view(out.size(0), -1)
        out = self.fc1(out)
        out = self.relu(out)
        out = self.fc2(out)
        out = self.sigmoid(out)
        out = out.view(out.size(0), out.size(1), 1, 1)
        out = out * original_out

        out += residual
        out = self.relu(out)

        return out


class SE_ResNeXt(nn.Module):

    def __init__(self, block, layers, num_classes=1000, num_group=32):
        self.inplanes = 64
        super(SE_ResNeXt, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0], num_group)
        self.layer2 = self._make_layer(block, 128, layers[1], num_group, stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], num_group, stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], num_group, stride=2)
        self.avgpool = nn.AvgPool2d(7, stride=1)
        self.fc = nn.Linear(512 * block.expansion, num_classes)

        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))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, blocks, num_group, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample, num_group=num_group))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes, num_group=num_group))

        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)

        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)

        return x

def se_resnext_50(**kwargs):
    """Constructs a ResNeXt-50 model.
    """
    model = SE_ResNeXt(Bottleneck, [3, 4, 6, 3], **kwargs)
    return model


def se_resnext_101(**kwargs):
    """Constructs a ResNeXt-101 model.
    """
    model = SE_ResNeXt(Bottleneck, [3, 4, 23, 3], **kwargs)
    return model


def se_resnext_152(**kwargs):
    """Constructs a ResNeXt-152 model.
    """
    model = SE_ResNeXt(Bottleneck, [3, 8, 36, 3], **kwargs)
    return model