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)))