Pytorch 目标分类比赛入门

视觉盛宴

已于 2022-10-29 15:26:08 修改

阅读量1.6k

点赞数 1

文章标签： pytorch 分类机器学习神经网络人工智能

于 2022-01-08 16:54:51 首次发布

本文链接：https://blog.csdn.net/cp1314971/article/details/122377437

版权

本文介绍了深度学习入门级别的图像分类比赛流程，包括数据加载、模型构建、训练和预测。使用了timm库，以102种鲜花分类为例，详细解释了数据预处理、ResNet18模型的搭建、训练过程以及如何生成预测结果。同时，文章还探讨了数据增强、学习率调整等技巧，并提供了加载权重继续训练的方法。

摘要由CSDN通过智能技术生成

分类比赛叙述

目标分类比赛是入门比较简单和编写代码量较少，适合初学者第一步比赛入门的基础。比赛一般分为初赛复赛，初赛大部分提交csv文件，复赛可能会需要在官方指定的服务器上docker环境部署然后进行预测服务器上测试集照片。

比赛数据

比赛连接AI研习社非常适合入门比赛，里面有很多不同的比赛和代码分享，大家可以去学习学习。数据下载我就以这个比赛的102鲜花分类为例。讲解比赛流程

数据加载

在这里插入图片描述
数据是这个样子的，train是我们要训练的数据图片，train.csv是保存的对应train下照片的label，而test只有照片没有label。我们目的是用train的数据训练模型然后对test进行预测生成test.csv然后提交成绩就可以了。

train_df = pd.read_csv('./54_data/train.csv')
train_df['filename'] = train_df['filename'].apply(lambda x: './54_data/train/{0}'.format(x))

class MyDataset(Dataset):
    def __init__(self, df, transform):
        self.df = df
        self.transform = transform

    def __getitem__(self, index):
        img = Image.open(self.df['filename'].iloc[index]).convert('RGB')
        img = self.transform(img)
        return img, torch.from_numpy(np.array(self.df['label'].iloc[index]))

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


train_transform = transforms.Compose([
    transforms.RandomRotation(15),
    transforms.Resize([300, 300]),
    transforms.RandomVerticalFlip(),
    #FixedRotation([0, 90, 180, -90]),
    transforms.RandomRotation(90),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
val_transform = transforms.Compose([
    transforms.Resize([300, 300]),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

train_data = MyDataset(train_df, train_transform)
trainloader = DataLoader(train_data, batch_size=16 , shuffle=True)

这一部分代码是读取数据然后整理成迭代器的方式。方便数据处理、里面有很多细节，比如class MyDataset这个类就是重写类。这里我只大概讲解比赛流程。详细部分我博客其他部分讲解的有。数据详解

模型构建

比赛当然是使用各种库。比如timm库或者mmclass库，这里我们使用timm库。而却就使用timm库中的模型。其他版块代码想要使用直接copy到我们代码里面就行了。就可以很方便的调用。

import timm
net=timm.create_model(“resnet18”,pretrained=True,num_classes=102).cuda()


import timm
model_list=timm.list_models("eff*",pretrained=True)
print(model_list)

就这简单的一行代码就构建好模型了。可以使用timm.list_models这个函数查看全部的模型。也可以使用正则表达式的方式搜索一些指定模型进行选择。

训练部分

model_name="efficientnet_b2"
epoch_num=22
net=timm.create_model(model_name,pretrained=True,num_classes=102).cuda()
train_data = MyDataset(train_df, train_transform)
trainloader = DataLoader(train_data, batch_size=16 , shuffle=True)
criterion=nn.CrossEntropyLoss()

optimizer = optim.SGD(net.parameters(),lr=0.02)
#optimizer = torch.optim.Adam(net.parameters(), lr=0.0000001)
#lr_scheduler = CosineLRScheduler(optimizer, t_initial=0.02, lr_min=0.000004)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
PATH = './dir/'+model_name
pre_acc = 0
def get_cur_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']
if __name__ == '__main__':
    correct = 0
    total = 0
    for epoch in range(epoch_num):
        print("========== epoch: [{}/{}] ==========".format(epoch + 1, epoch_num))
        for i, (inputs, labels) in tqdm(enumerate(trainloader)):
            inputs = inputs.cuda()
            labels = labels.cuda()
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            correct += (outputs.argmax(dim=1) == labels).sum().item()
            total += labels.size(0)
            train_acc = 100.0 * correct / total
            optimizer.zero_grad()  # 梯度先全部降为0
            loss.backward()  # 反向传递过程
            optimizer.step()  # 以学习效率0.001来优化梯度
            if i % 10 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f} train acc: {:.6f}  lr : {:.6f}'.format(epoch+1, i * len(inputs),
                len(trainloader.dataset),100. * i / len(trainloader),loss.item(),train_acc,get_cur_lr(optimizer)))
        if epoch % 1 ==0:
            torch.save({'epoch': epoch,
                        'model_state_dict': net.state_dict(),
                           'optimizer_state_dict': optimizer.state_dict(),
                        'loss': loss
            },PATH+"_"+str(epoch+1)+".pth")
        
        lr_scheduler.step()

以上代码属于基本功。就是定义好网络，我们选择优化器，学习率等等。然后保存权重

全部代码

import torchvision.transforms as transforms
import torch.nn as nn
import numpy as np
import pandas as pd
import torch.optim as optim
from PIL import Image
from tqdm import tqdm
import torch.nn.functional as F
import torch
import random
from torch.optim import lr_scheduler
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
import timm
from timm.scheduler.cosine_lr import CosineLRScheduler
train_df = pd.read_csv('./54_data/train.csv')  #这里需要注意一下你数据集的位置
#print(train_df)
#print(train_df['filename'][0])
train_df['filename'] = train_df['filename'].apply(lambda x: './54_data/train/{0}'.format(x))
# val_df = pd.read_csv('./54_data/val.csv')
# val_df['filename'] = val_df['filename'].apply(lambda x: './54_data/train/{0}'.format(x))


class MyDataset(Dataset):
    def __init__(self, df, transform):
        self.df = df
        self.transform = transform

    def __getitem__(self, index):
        img = Image.open(self.df['filename'].iloc[index]).convert('RGB')
        img = self.transform(img)
        return img, torch.from_numpy(np.array(self.df['label'].iloc[index]))

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


train_transform = transforms.Compose([
    transforms.RandomRotation(15),
    transforms.Resize([300, 300]),
    transforms.RandomVerticalFlip(),
    #FixedRotation([0, 90, 180, -90]),
    transforms.RandomRotation(90),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
val_transform = transforms.Compose([
    transforms.Resize([300, 300]),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
model_name="efficientnet_b2"
epoch_num=22
net=timm.create_model(model_name,pretrained=True,num_classes=102).cuda()
train_data = MyDataset(train_df, train_transform)
trainloader = DataLoader(train_data, batch_size=16 , shuffle=True)
criterion=nn.CrossEntropyLoss()

optimizer = optim.SGD(net.parameters(),lr=0.02)
#optimizer = torch.optim.Adam(net.parameters(), lr=0.0000001)
#lr_scheduler = CosineLRScheduler(optimizer, t_initial=0.02, lr_min=0.000004)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
PATH = './'+model_name  #这里也需要注意，你需要先新建一个dir文件夹，一会在这个文件下存放权重
pre_acc = 0
def get_cur_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']
if __name__ == '__main__':
    correct = 0
    total = 0
    for epoch in range(epoch_num):
        print("========== epoch: [{}/{}] ==========".format(epoch + 1, epoch_num))
        for i, (inputs, labels) in tqdm(enumerate(trainloader)):
            inputs = inputs.cuda()
            labels = labels.cuda()
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            correct += (outputs.argmax(dim=1) == labels).sum().item()
            total += labels.size(0)
            train_acc = 100.0 * correct / total
            optimizer.zero_grad()  
            loss.backward()  
            optimizer.step()  
            if i % 10 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f} train acc: {:.6f}  lr : {:.6f}'.format(epoch+1, i * len(inputs),
                len(trainloader.dataset),100. * i / len(trainloader),loss.item(),train_acc,get_cur_lr(optimizer)))
        if epoch % 1 ==0:
            torch.save({'epoch': epoch,
                        'model_state_dict': net.state_dict(),
                           'optimizer_state_dict': optimizer.state_dict(),
                        'loss': loss
            },PATH+"_"+str(epoch+1)+".pth")
        
        lr_scheduler.step()

想要成功运行这部分代码，需要处理好数据集的位置和你需要提前安装好各种包。比如timm框架之类的。然后就可以成功运行训练了。
在这里插入图片描述
时间关系我就简单训练5个批次吧，单模还是很容易到98分的成绩训练五个模型差不多集成就可以100分貌似。

预测结果

 	import torch.optim as optim
import torchvision.transforms as transforms
import tqdm
import torch
import os
import numpy as np
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import timm

model_name="efficientnet_b2"
net=timm.create_model(model_name,pretrained=False,num_classes=102).cuda()
PATH = './efficientnet_b2_5.pth'
checkpoint = torch.load(PATH)
net.load_state_dict(checkpoint['model_state_dict'])
test='./54_data/test'
data_len = len(os.listdir('./54_data/test'))
test_path_list = ['{}/{}.jpg'.format(test, x) for x in range(0, data_len)]
test_data = np.array(test_path_list)

class MyDataset(Dataset):
    def __init__(self, df, transform, mode='train'):
        self.df = df
        self.transform = transform
        self.mode = mode

    def __getitem__(self, index):
        if self.mode == 'train':
            img = Image.open(self.df['filename'].iloc[index]).convert('RGB')
            img = self.transform(img)
            return img, torch.from_numpy(np.array(self.df['label'].iloc[index]))
        else:
            img = Image.open(self.df[index]).convert('RGB')
            img = self.transform(img)
            return img, torch.from_numpy(np.array(0))

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

test_transform = transforms.Compose([
    transforms.Resize([300, 300]),
    transforms.ToTensor(),
])

test_dataset = MyDataset(test_data, test_transform,'test')
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
net.eval()
pred_list = []
with torch.no_grad():
    for batch_x, (img,label) in enumerate(test_loader):
        print(batch_x)
        img = img.cuda()
        # compute output
        probs = net(img)
        preds = torch.argmax(probs, dim=1)
        pred_list += [p.item() for p in preds]
submission = pd.DataFrame({"id": range(len(pred_list)), "label": pred_list})
submission.to_csv('submission.csv', index=False, header=False)

提交成绩

比赛连接
在这里插入图片描述

这仅仅训练五个批次。到此就走了一遍比赛的流程。

Trick修改

auto_augment

timm里面比赛常见的tirck也都有。比如auto增强，mixup，cutmix都已经写好了，我们直接拿过来用就ok了，不同的trick在不同的模型涨分情况不一样，是需要大量尝试的，代码如下

from timm.data.auto_augment import auto_augment_transform
tfm = auto_augment_transform(config_str = 'original', 
hparams = {'translate_const': 
100, 'img_mean': (124, 116, 104)})

这里引用好。然后我们在封装成一个类，就可以直接加入到transform里面

class auto(object):
    def __init__(self, aa=1):
        self.aa = aa
    def __call__(self, img):
        return tfm(img)

写成这种类，self.aa没啥含义，只是格式。之后我们就可以直接使用了

train_transform = transforms.Compose([                               
    auto(),
    transforms.Resize([300, 300]),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

这样就算成功使用了。

任意数据增强

class FixedRotation(object):
    def __init__(self, angles):
        self.angles = angles

    def __call__(self, img):
        return fixed_rotate(img, self.angles)


def fixed_rotate(img, angles):
    angles = list(angles)
    angles_num = len(angles)
    index = random.randint(0, angles_num - 1)
    return img.rotate(angles[index])


train_transform = transforms.Compose([
    transforms.RandomRotation(15),
    transforms.Resize([300, 300]),
    transforms.RandomVerticalFlip(),
    FixedRotation([0, 90, 180, -90]),
    transforms.RandomRotation(90),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

配合一些库写好现成的，我们只需要稍微修改一下封装成类即可。然后在Compose直接使用。这样就可以写出任意自己想要的效果。

自定义学习率

def adjust_learning_rate(optimizer, epoch):
    global lr

    if epoch <= 5:
        lr = 0.001*epoch
    elif epoch<= 15:
            lr = 0.02
    elif epoch <=20:
            lr = 0.002
    elif epoch <=25:
            lr = 0.0002
    elif epoch <=35:
            lr =0.00002
    else:
        lr = 0.000001

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

写成一个函数。然后在train函数中调用传入参数就可以了

lr_scheduler = adjust_learning_rate(optimizer, epoch)
 #lr_scheduler.step()

这里需要注意一点就是需要把step()这个函数注释掉。我们已经在函数里面定义修改lr的值了所以不需要在自动更新lr的值。自定义学习率是通过训练一次模型观察acc的变化，是靠经验来学习设置了。比如当你发现学习率处于0.0002，acc很难变动了，在训练几个批次也无济于事，那么这个时候就应该适当缩小学习率，避免不不需要的训练批次。浪费显卡资源。

加载权重，继续训练

def resume(PATH):
  checkpoint = torch.load(PATH)
  net.load_state_dict(checkpoint['model_state_dict'])
  optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
  epoch = checkpoint['epoch']
  loss = checkpoint['loss']

一个简单的函数就可以搞定，这样就可以接着上次中断的权重继续训练了。

其他比赛训练格式

数据加载第二种训练方式

如果不用csv文件，需要确保数据在train文件夹下按照每个类别。如下图所示，每个类别下全是jpg或者其他格式图像
在这里插入图片描述
训练代码

from torchvision.datasets import ImageFolder
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms
import torch
import torchvision.transforms as transforms
import torch.nn as nn
import numpy as np
import pandas as pd
import torch.optim as optim
from PIL import Image
from tqdm import tqdm
import torch.nn.functional as F
import torch
import random
from torch.optim import lr_scheduler
from torch.autograd import Variable
import timm
from timm.scheduler.cosine_lr import CosineLRScheduler
from utils_aug import ImageNetPolicy
train_path= './datasets/train/'
test_path='./datasets/test/'
train_transform=transforms.Compose([
    #transforms.RandomRotation(15),
    transforms.Resize([384, 384]),
    #transforms.RandomVerticalFlip(),
    #transforms.RandomRotation(90),
    ImageNetPolicy(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

test_transform=transforms.Compose([
    transforms.Resize([384, 384]),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])


train_dataset=ImageFolder(train_path,transform=train_transform)
test_dataset=ImageFolder(test_path,transform=test_transform)
train_trainloader = torch.utils.data.DataLoader(train_dataset, batch_size=12,shuffle=True, num_workers=0)
test_trainloader = torch.utils.data.DataLoader(test_dataset, batch_size=12,shuffle=True, num_workers=0)
model_name="convnext_xlarge_in22k"
epoch_num=28
net=timm.create_model(model_name,pretrained=True,num_classes=18).cuda()
criterion=nn.CrossEntropyLoss()

optimizer = optim.SGD(net.parameters(),lr=0.02)
#optimizer = torch.optim.Adam(net.parameters(), lr=0.0000001)
#lr_scheduler = CosineLRScheduler(optimizer, t_initial=0.02, lr_min=0.000004)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=12, gamma=0.1)
PATH = './'+model_name  #这里也需要注意，你需要先新建一个dir文件夹，一会在这个文件下存放权重
pre_acc = 0


def get_cur_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']


if __name__ == '__main__':
    correct = 0
    total = 0
    for epoch in range(epoch_num):
        print("========== epoch: [{}/{}] ==========".format(epoch + 1, epoch_num))
        for i, (inputs, labels) in tqdm(enumerate(train_trainloader)):
            inputs = inputs.cuda()
            labels = labels.cuda()
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            correct += (outputs.argmax(dim=1) == labels).sum().item()
            total += labels.size(0)
            train_acc = 100.0 * correct / total
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if i % 10 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f} train acc: {:.6f}  lr : {:.6f}'.format(epoch + 1,
                                                                                                              i * len(
                                                                                                                  inputs),
                                                                                                              len(train_trainloader.dataset),
                                                                                                              100. * i / len(
                                                                                                                  train_trainloader),
                                                                                                              loss.item(),
                                                                                                              train_acc,
                                                                                                              get_cur_lr(
                                                                                                                  optimizer)))
        if epoch+1 == epoch_num:
            torch.save({'epoch': epoch,
                        'model_state_dict': net.state_dict(),
                        'optimizer_state_dict': optimizer.state_dict(),
                        'loss': loss
                        }, PATH + "_" + str(epoch + 1) + ".pth")

        lr_scheduler.step()
        total, correct = 0, 0 
        net.eval()

        with torch.no_grad():
            print("*************** test ***************")
            for X, y in test_trainloader:
                X, y = X.to('cuda'), y.to('cuda')

                output = net(X)
                loss = criterion(output, y)

                total += y.size(0)
                correct += (output.argmax(dim=1) == y).sum().item()

        test_acc = 100.0 * correct / total

        print("test_loss: {:.3f} | test_acc: {:6.3f}%" \
              .format(loss.item(), test_acc))
        print("************************************\n")
        net.train()

albu数据增强的训练代码

import torchvision.transforms as transforms
import torch.nn as nn
import numpy as np
import torch.optim as optim
from torch.utils.data.sampler import SubsetRandomSampler
from PIL import Image
from tqdm import tqdm
import torch.nn.functional as F
import torch
import random
import cv2
import timm
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
import albumentations as A
from albumentations.pytorch import ToTensorV2
train_df = pd.read_csv('./all.csv')
train_df['path'] = train_df['path'].apply(lambda x: './img/{0}'.format(x))

class FocalLoss(nn.Module):
    def __init__(self, gamma=0, alpha=None, size_average=True):
        super(FocalLoss, self).__init__()
        self.gamma = gamma
        self.alpha = alpha
        if isinstance(alpha, (float, int)): self.alpha = torch.Tensor([alpha, 1 - alpha])
        if isinstance(alpha, list): self.alpha = torch.Tensor(alpha)
        self.size_average = size_average

    def forward(self, input, target):
        if input.dim() > 2:
            input = input.view(input.size(0), input.size(1), -1)  # N,C,H,W => N,C,H*W
            input = input.transpose(1, 2)  # N,C,H*W => N,H*W,C
            input = input.contiguous().view(-1, input.size(2))  # N,H*W,C => N*H*W,C
        target = target.view(-1, 1)
        logpt = F.log_softmax(input, dim=1)
        logpt = logpt.gather(1, target)
        logpt = logpt.view(-1)
        pt = Variable(logpt.data.exp())

        if self.alpha is not None:
            if self.alpha.type() != input.data.type():
                self.alpha = self.alpha.type_as(input.data)
            at = self.alpha.gather(0, target.data.view(-1))
            logpt = logpt * Variable(at)

        loss = -1 * (1 - pt) ** self.gamma * logpt
        if self.size_average:
            return loss.mean()
        else:
            return loss.sum()

class FixedRotation(object):
    def __init__(self, angles):
        self.angles = angles

    def __call__(self, img):
        return fixed_rotate(img, self.angles)


def fixed_rotate(img, angles):
    angles = list(angles)
    angles_num = len(angles)
    index = random.randint(0, angles_num - 1)
    return img.rotate(angles[index])


class MyDataset(Dataset):
    def __init__(self, df, transform):
        self.df = df
        self.transform = transform

    def __getitem__(self, index):
        #img = Image.open(self.df['path'].iloc[index]).convert('RGB')
        img = cv2.imread(self.df['path'].iloc[index])
        #print(self.df['path'].iloc[index])
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        #img = self.transform(img)
        if self.transform is not None:
          img = self.transform(image=img)["image"]
        return img, torch.from_numpy(np.array(self.df['label'].iloc[index]))

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

train_transform2 = A.Compose([
          A.Resize(384,384),
          A.HorizontalFlip(p=0.5),
          #A.HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.5),
          #A.RandomBrightnessContrast(brightness_limit=(-0.1,0.1), contrast_limit=(-0.1, 0.1), p=0.5),
          A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),
          A.CoarseDropout(p=0.5),
          A.Cutout(p=0.5),
          ToTensorV2()       
])



model_name='convnext_tiny'
train_data = MyDataset(train_df, train_transform2)
trainloader = DataLoader(train_data, batch_size=16 , shuffle=True)
net=timm.create_model(model_name,pretrained=True,num_classes=20).cuda()
criterion=FocalLoss(0.5)
#criterion=nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(),lr=0.02)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
pth_PATH='./weights_conv_img384/'
!mkdir weights_conv_img384
epoch_num= 25
correct = 0
total = 0





def get_cur_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']
for epoch in range(epoch_num):
  print("========== epoch: [{}/{}] ==========".format(epoch + 1, epoch_num))
  for i, (inputs, labels) in tqdm(enumerate(trainloader)):
    inputs = inputs.cuda()
    labels = labels.cuda()
    outputs = net(inputs)
    loss = criterion(outputs, labels)
    correct += (outputs.argmax(dim=1) == labels).sum().item()
    total += labels.size(0)
    train_acc = 100.0 * correct / total
    optimizer.zero_grad()  # 梯度先全部降为0
    loss.backward()  # 反向传递过程
    optimizer.step()  # 以学习效率0.001来优化梯度
    if i % 10 == 0:
      print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f} train acc: {:.6f}  lr : {:.6f}'.format(epoch+1, i * len(inputs),
      len(trainloader.dataset),100. * i / len(trainloader),loss.item(),train_acc,get_cur_lr(optimizer)))
  if (epoch+1)  == epoch_num or (epoch%10==0) :
      torch.save({'epoch': epoch,
                  'model_state_dict': net.state_dict(),
                      'optimizer_state_dict': optimizer.state_dict(),
                  'loss': loss
      },pth_PATH+model_name+"_"+str(epoch+1)+"_"+"224"+".pth")

  lr_scheduler.step()

推理代码，单张推理

import torch.optim as optim
import torchvision.transforms as transforms
from torchvision import datasets, transforms
import tqdm
import torch
import os
import numpy as np
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import timm
from torch.autograd import Variable

model_name="convnext_xlarge_in22k"
net=timm.create_model(model_name,pretrained=False,num_classes=18).cuda()
PATH = './convnext_xlarge_in22k_50.pth'
checkpoint = torch.load(PATH)
net.load_state_dict(checkpoint['model_state_dict'])
test='./SAR_test_nolabel/'
data = os.listdir('./SAR_test_nolabel/')
#test_path_list = ['{}/{}.bmp'.format(test, x) for x in range(0, data_len)]
#test_data = np.array(test_path_list)
test_transform=transforms.Compose([
    transforms.Resize([224, 224]),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
n=0
net.eval()
pred_list=[]
with torch.no_grad():
    for img_name in data:
        tt=''
        n=n+1
        img=Image.open(test+img_name).convert("RGB")
        
        img=test_transform(img)
        img = Variable(torch.unsqueeze(img,dim=0).float(),requires_grad=False)
        img = img.cuda()
        # compute output
        probs = net(img)
        preds = torch.argmax(probs, dim=1)
        name=img_name[:-4]
        tt=name+" "+str(int(preds))+"/n"
        print(tt)
        pred_list.append(tt)
with open("./test.txt","w") as f:
     f.write(str(pred_list))