Pytorch ShuffeleNetV2回归图像

训练

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.backends.cudnn as cudnn
import numpy as np
from data_gen import OccGenDataset
from model import ShuffleNetV2
import os
import argparse
parser = argparse.ArgumentParser(description='Traing Arguements')
parser.add_argument('--output', '-o', default='result', type=str, help='Traing result')
parser.add_argument('--lr', default=0.01, type=float, help='lr')
parser.add_argument('--cuda', default='True', type=str, help='cuda or cpu')
parser.add_argument('--batch_size', default=32, type=int, help='batch_size')
parser.add_argument('--gpu', '-g', default='0', type=str, help='gpu id')
parser.add_argument('--step_value', default=[60, 80], type=list, help='lr decay step')
parser.add_argument('--max_epoch', default=100, type=int, help='training epochs')
args = parser.parse_args()

os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
device = 'cuda' if args.cuda else 'cpu'

if not os.path.exists(args.output):
    os.mkdir(args.output)



def train(train_loader, model, criterion, optimizer, epoch):
    model.train()
    train_loss = 0
    total = 0
    for batch_idx, (inputs, targets) in enumerate(train_loader):
        inputs, targets = inputs.to(device), targets.float().to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        outputs = torch.squeeze(outputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        
        train_loss += loss.item()
        total += targets.size(0)
        MAE = np.mean(np.abs(outputs.cpu().data.numpy(), targets.cpu().data.numpy()))
        
        if batch_idx % 100 == 0:
            print('Train Epoch:{}({}/{}) loss:{:.4f} MAE:{:.4f}'.format(
                epoch, batch_idx, len(train_loader), 
                train_loss/(batch_idx+1), MAE/total*100))
            f.write('Train Epoch:{}({}/{}) loss:{:.4f} MAE:{:.4f}\n'.format(
                epoch, batch_idx, len(train_loader), 
                train_loss/(batch_idx+1), MAE/total*100))
            

def vaild(vaild_loader, model, criterion, optimizer, epoch):
    model.eval()
    val_loss = 0
    total = 0
    with torch.no_grad():
        for batch_idx, (inputs, targets) in enumerate(vaild_loader):
            inputs, targets = inputs.to(device), targets.float().to(device)
            outputs = model(inputs)
            outputs = torch.squeeze(outputs)
            loss = criterion(outputs, targets)
            val_loss += loss.item()
            total += targets.size(0)
            MAE = np.mean(np.abs(outputs.cpu().data.numpy(), targets.cpu().data.numpy()))

            if batch_idx % 100 == 0:
                print('Val Epoch:{}({}/{}) loss:{:.4f} MAE:{:.4f}'.format(
                    epoch, batch_idx, len(vaild_loader), 
                    val_loss/(batch_idx+1), MAE/total*100))
                f.write('Val Epoch:{}({}/{}) loss:{:.4f} MAE:{:.4f}\n'.format(
                    epoch, batch_idx, len(vaild_loader), 
                    val_loss/(batch_idx+1), MAE/total*100))


def adjust_opt(optimizer, gamma, step_index, epoch):
    lr = args.lr * (gamma**step_index)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    return lr


def train_net(dataset_loader, model, criterion, optimizer, max_epoch, step_value):
    train_loader, vaild_loader = dataset_loader
    step_index = 0
    for epoch in range(0, max_epoch):
        if epoch in step_value:
            step_index += 1
        lr = adjust_opt(optimizer, 0.1, step_index, epoch)
        train(train_loader, model, criterion, optimizer, epoch)
        vaild(vaild_loader, model, criterion, optimizer, epoch)
        if epoch%10 == 0:
            torch.save(model.state_dict(), os.path.join(args.output,'model_{}.pt'.format(epoch)))
    

if __name__ == '__main__':
    train_dataset = OccGenDataset('train', '../images', './occ_train.txt')
    vaild_dataset = OccGenDataset('valid', '../images', './occ_val.txt')
    train_loader = torch.utils.data.DataLoader(train_dataset, 
                     batch_size=args.batch_size, shuffle=True, num_workers=2)
    vaild_loader = torch.utils.data.DataLoader(vaild_dataset, 
                      batch_size=args.batch_size, shuffle=True, num_workers=2)
    dataset_loader = (train_loader, vaild_loader)
    model = ShuffleNetV2(1, 1).to(device)
    criterion = nn.L1Loss().to(device)
    #criterion = nn.MSELoss()
    optimizer =  optim.SGD(model.parameters(), lr=args.lr, 
                      momentum=0.9, weight_decay=5e-4)
    f = open('log.txt', 'w')
    train_net(dataset_loader, model, criterion, optimizer, 
                      args.max_epoch, args.step_value)
    f.close()

Demo

import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import cv2
import numpy as np
from model import ShuffleNetV2
import os
import argparse
from torchvision import transforms
from PIL import Image
parser = argparse.ArgumentParser(description='Traing Arguements')
parser.add_argument('--cuda', default='True', type=str, help='cuda or cpu')
parser.add_argument('--image', default='./demo.jpg', type=str, help='iamge')
parser.add_argument('--gpu', '-g', default='0', type=str, help='gpu id')
parser.add_argument('--model', default='./result_64/model_0.pt', type=str, help='gpu id')
args = parser.parse_args()

data_transforms = {
    'valid': transforms.Compose([
        transforms.Resize((64, 64)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}



if __name__ == '__main__':
    model = ShuffleNetV2(1, 1)
    model.load_state_dict(torch.load(args.model))
    img = Image.open('../images/mask/1521761533859.jpg')
    model.eval()
    img = data_transforms['valid'](img)
    img = img.unsqueeze(0)
    num = 1000
    import time 
    tic = time.time()
    for _ in range(num):
        tic1 = time.time()
        out = model(img)
        print(time.time()-tic1)
    out = out.squeeze(0)
    print('score:{} time:{}'.format(out.cpu().data.numpy()[0]*100, 
           (time.time()-tic)/num))

网络结构（input: (3x64x64)）

'''ShuffleNetV2 in PyTorch.

See the paper "ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design" for more details.
'''
import torch
import torch.nn as nn
import torch.nn.functional as F


class ShuffleBlock(nn.Module):
    def __init__(self, groups=2):
        super(ShuffleBlock, self).__init__()
        self.groups = groups

    def forward(self, x):
        '''Channel shuffle: [N,C,H,W] -> [N,g,C/g,H,W] -> [N,C/g,g,H,w] -> [N,C,H,W]'''
        N, C, H, W = x.size()
        g = self.groups
        return x.view(N, g, int(C/g), H, W).permute(0, 2, 1, 3, 4).reshape(N, C, H, W)


class SplitBlock(nn.Module):
    def __init__(self, ratio):
        super(SplitBlock, self).__init__()
        self.ratio = ratio

    def forward(self, x):
        c = int(x.size(1) * self.ratio)
        return x[:, :c, :, :], x[:, c:, :, :]


class BasicBlock(nn.Module):
    def __init__(self, in_channels, split_ratio=0.5):
        super(BasicBlock, self).__init__()
        self.split = SplitBlock(split_ratio)
        in_channels = int(in_channels * split_ratio)
        self.conv1 = nn.Conv2d(in_channels, in_channels,
                               kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(in_channels)
        self.conv2 = nn.Conv2d(in_channels, in_channels,
                               kernel_size=3, stride=1, padding=1, groups=in_channels, bias=False)
        self.bn2 = nn.BatchNorm2d(in_channels)
        self.conv3 = nn.Conv2d(in_channels, in_channels,
                               kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(in_channels)
        self.shuffle = ShuffleBlock()

    def forward(self, x):
        x1, x2 = self.split(x)
        out = F.relu(self.bn1(self.conv1(x2)))
        out = self.bn2(self.conv2(out))
        out = F.relu(self.bn3(self.conv3(out)))
        out = torch.cat([x1, out], 1)
        out = self.shuffle(out)
        return out


class DownBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DownBlock, self).__init__()
        mid_channels = out_channels // 2
        # left
        self.conv1 = nn.Conv2d(in_channels, in_channels,
                               kernel_size=3, stride=2, padding=1, groups=in_channels, bias=False)
        self.bn1 = nn.BatchNorm2d(in_channels)
        self.conv2 = nn.Conv2d(in_channels, mid_channels,
                               kernel_size=1, bias=False)
        self.bn2 = nn.BatchNorm2d(mid_channels)
        # right
        self.conv3 = nn.Conv2d(in_channels, mid_channels,
                               kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(mid_channels)
        self.conv4 = nn.Conv2d(mid_channels, mid_channels,
                               kernel_size=3, stride=2, padding=1, groups=mid_channels, bias=False)
        self.bn4 = nn.BatchNorm2d(mid_channels)
        self.conv5 = nn.Conv2d(mid_channels, mid_channels,
                               kernel_size=1, bias=False)
        self.bn5 = nn.BatchNorm2d(mid_channels)

        self.shuffle = ShuffleBlock()

    def forward(self, x):
        # left
        out1 = self.bn1(self.conv1(x))
        out1 = F.relu(self.bn2(self.conv2(out1)))
        # right
        out2 = F.relu(self.bn3(self.conv3(x)))
        out2 = self.bn4(self.conv4(out2))
        out2 = F.relu(self.bn5(self.conv5(out2)))
        # concat
        out = torch.cat([out1, out2], 1)
        out = self.shuffle(out)
        return out


class ShuffleNetV2(nn.Module):
    def __init__(self, net_size=1, num_classes=10):
        super(ShuffleNetV2, self).__init__()
        out_channels = configs[net_size]['out_channels']
        num_blocks = configs[net_size]['num_blocks']
        self.num_classes = num_classes

        self.conv1 = nn.Conv2d(3, 24, kernel_size=3,
                               stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(24)
        self.in_channels = 24
        self.layer1 = self._make_layer(out_channels[0], num_blocks[0])
        self.layer2 = self._make_layer(out_channels[1], num_blocks[1])
        self.layer3 = self._make_layer(out_channels[2], num_blocks[2])
        self.conv2 = nn.Conv2d(out_channels[2], out_channels[3],
                               kernel_size=1, stride=1, padding=0, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels[3])
        self.linear = nn.Linear(out_channels[3], self.num_classes)

    def _make_layer(self, out_channels, num_blocks):
        layers = [DownBlock(self.in_channels, out_channels)]
        for i in range(num_blocks):
            layers.append(BasicBlock(out_channels))
            self.in_channels = out_channels
        return nn.Sequential(*layers)

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x))) # 3x64x64 -> 24x32x32
        out = F.max_pool2d(out, 3, stride=2, padding=1) # ->24x16x16
        out = self.layer1(out) # -> 116x8x8
        out = self.layer2(out) # -> 232x4x4
        out = self.layer3(out) # -> 464x2x2
        out = F.relu(self.bn2(self.conv2(out)))
        out = F.avg_pool2d(out, 2)   # Nx1024x1x1
        out = out.view(out.size(0), -1) # Nx1024
        out = self.linear(out)       # Nx1
        out = F.sigmoid(out)
        return out


configs = {
    0.5: {
        'out_channels': (48, 96, 192, 1024),
        'num_blocks': (3, 7, 3)
    },

    1: {
        'out_channels': (116, 232, 464, 1024),
        'num_blocks': (3, 7, 3)
    },
    1.5: {
        'out_channels': (176, 352, 704, 1024),
        'num_blocks': (3, 7, 3)
    },
    2: {
        'out_channels': (224, 488, 976, 2048),
        'num_blocks': (3, 7, 3)
    }
}



if __name__ == '__main__':
    net = ShuffleNetV2(0.5, 10)
    from torchsummary import summary
    net = net.to('cuda')
    summary(net, (3, 224, 224))