mobilenetv2图像分类

from torch.utils import data
import os
import cv2
from torchvision import transforms as T
from torchvision.models.mobilenetv2 import mobilenet_v2
from torch.optim.lr_scheduler import StepLR
from torch.nn import CrossEntropyLoss  # SmoothL1Loss
import torch
class datasest(data.Dataset):
    def __init__(self, path):
        self.img_path = []
        self.labels = []
        self.transforms = T.Compose([
            lambda x: cv2.resize(x, (224, 224)),
            T.ToTensor(),
            T.Normalize(mean=[0.5], std=[0.5]),
        ])
        for dir in os.listdir(path):
            for file in os.listdir(os.path.join(path, dir)):
                self.img_path.append(os.path.join(path, dir, file))
                self.labels.append(int(dir))

    def __len__(self):
        return len(self.img_path)

    def __getitem__(self, index):
        img = cv2.imread(self.img_path[index])
        img = self.transforms(img)
        return img, self.labels[index]


train_root = r'D:\test\ceshi\fenlei'
batch_size = 30
lr = 1e-3
weight_decay = 1e-5
lr_step = 10
lr_decay = 0.5  # when val_loss increase, lr = lr*lr_decay
max_epoch = 1000

model = mobilenet_v2(pretrained=False)
traindataset = datasest(train_root)
trainloader = torch.utils.data.DataLoader(traindataset, batch_size=batch_size, shuffle=True, num_workers=0)

criterion = CrossEntropyLoss()  # SmoothL1Loss()

model.cuda()

optimizer = torch.optim.Adam([{'params': model.parameters()}], lr=lr,
                             weight_decay=weight_decay)

scheduler = StepLR(optimizer, step_size=lr_step, gamma=lr_decay)
for epoch in range(max_epoch):
    model.train()
    for batch, (data_input, label) in enumerate(trainloader):
        data_input = data_input.cuda()
        label = label.cuda()
        output = model(data_input)
        loss = criterion(output, label)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        iters = epoch * len(trainloader) + batch
        if iters % 5 == 0:
            print(loss)
    if epoch % 5 == 0:
        torch.save(model.state_dict(), 'net.pth')

预测：

import cv2
from torchvision import transforms as T
from torchvision.models.mobilenetv2 import mobilenet_v2
import torch

transforms = T.Compose([
            T.ToTensor(),
            T.Normalize(mean=[0.5], std=[0.5]),
        ])

model = mobilenet_v2()
model.load_state_dict(torch.load('net.pth'))
model.eval()
img = cv2.imread(r'D:\test\ceshi\fenlei\7\504003322861042723050406261_UP2_15_1254.jpg')
img = cv2.resize(img, (224, 224))
img = transforms(img)
img = torch.reshape(img,(1,3,224,224))
print(img.shape)
print(torch.argmax(model(img)))