图像分类实战：mobilenetv2从训练到TensorRT部署（pytorch）

imgs.append(imgpath)

trainval_files, val_files = train_test_split(imgs, test_size=0.3, random_state=42)

if train:

self.imgs = trainval_files

else:

self.imgs = val_files

def getitem(self, index):

“”"

一次返回一张图片的数据

“”"

img_path = self.imgs[index]

img_path=img_path.replace(“\”,‘/’)

if self.test:

label = -1

else:

labelname = img_path.split(‘/’)[-2]

label = Labels[labelname]

data = Image.open(img_path).convert(‘RGB’)

data = self.transforms(data)

return data, label

def len(self):

return len(self.imgs)

说一下代码的核心逻辑：

第一步 建立字典，定义类别对应的ID,用数字代替类别。

第二步 在__init__里面编写获取图片路径的方法。测试集只有一层路径直接读取，训练集在train文件夹下面是类别文件夹，先获取到类别，再获取到具体的图片路径。然后使用sklearn中切分数据集的方法，按照7：3的比例切分训练集和验证集。

第三步 在__getitem__方法中定义读取单个图片和类别的方法，由于图像中有位深度32位的，所以我在读取图像的时候做了转换。

然后我们在train.py调用SeedlingData读取数据 ，记着导入刚才写的dataset.py(from dataset.dataset import SeedlingData)

dataset_train = SeedlingData(‘data/train’, transforms=transform, train=True)

dataset_test = SeedlingData(“data/train”, transforms=transform_test, train=False)

读取数据

print(dataset_train.imgs)

导入数据

train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)

设置模型

---

• 设置loss函数为nn.CrossEntropyLoss()。

• 设置模型为mobilenet_v3_large，预训练设置为true，num_classes设置为12。

• 优化器设置为adam。

• 学习率调整策略选择为余弦退火。

实例化模型并且移动到GPU

criterion = nn.CrossEntropyLoss()

#criterion = SoftTargetCrossEntropy()

model_ft = mobilenet_v2(pretrained=True)

print(model_ft)

num_ftrs = model_ft.classifier[1].in_features

model_ft.classifier[1] = nn.Linear(num_ftrs, 12,bias=True)

model_ft.to(DEVICE)

print(model_ft)

选择简单暴力的Adam优化器，学习率调低

optimizer = optim.Adam(model_ft.parameters(), lr=modellr)

cosine_schedule = optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer,T_max=20,eta_min=1e-9)

#model_ft=torch.load(“model_33_0.952.pth”)

定义训练和验证函数

---

定义训练过程

alpha=0.2

def train(model, device, train_loader, optimizer, epoch):

model.train()

sum_loss = 0

total_num = len(train_loader.dataset)

print(total_num, len(train_loader))

for batch_idx, (data, target) in enumerate(train_loader):

data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)

data, labels_a, labels_b, lam = mixup_data(data, target, alpha)

optimizer.zero_grad()

output = model(data)

loss = mixup_criterion(criterion, output, labels_a, labels_b, lam)

loss.backward()<