图片分类: 多类别

    最近需要训练一个有200多类的图片分类网络,搜了一遍,发现居然没有很合适用的开源项目,于是自己简单撸了一个轮子,项目地址: https://github.com/xuduo35/imgcls_pytorch。支持如下backbone:

  • alexnet
  • resnet18,resnet34,resnet50,resnet101, resnet152, resnext101_32x4d, resnext101_64x4d
  • vgg11_bn, vgg16_bn
  • densenet121, densenet169, densenet161
  • inceptionv3, inceptionv4, inceptionresnetv2, bninception
  • xception, xception_att
  • dpn98, dpn107, dpn131
  • senet154, se_resnet50, se_resnet101, se_resnet152, se_resnext50_32x4d
  • pnasnet5large
  • polynet
  • efficientnet

    使用简便,第一步是按如下格式准备数据集,

  • your_dataset_directory
    • class1
      • 1.jpg
      • 2.jpg
    • class2
      • 1.jpg
      • 2.jpg
      • ...
    • ...

    自定义一个Dataset,实现如下

class ImageFolderEx(Dataset):
    def __init__(self, image_dir, image_files, image_labels, classnum=1000, transform=None):
        self.image_dir = image_dir
        self.image_files = image_files
        self.image_labels = image_labels
        self.classnum = classnum
        self.transform = transform

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

    def __getitem__(self, index):
        image_name = os.path.join(self.image_dir, self.image_files[index])  
        image = cv2.imread(image_name)
        image = image[:,:,::-1]
        image = Image.fromarray(image)
        label = self.image_labels[index]
        if self.transform:
            image = self.transform(image)
        onehot = [0]*self.classnum
        onehot[label] = 1
        return (image, np.array(onehot).astype(np.float32))   

    支持简单的余弦退火学习率调度器

    scheduler = optim.lr_scheduler.CosineAnnealingLR(
            optimizer,
            T_max=args.epochs, eta_min=0.00001, last_epoch=-1
            )

    一方面执行简单的数据增广

            transforms.Compose([
                transforms.RandomRotation([-13,13]),
                transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5),
                transforms.Resize(args.imgsz+32),
                transforms.RandomCrop(args.imgsz),
                transforms.ToTensor(),
                normalize
            ])

    另外再按照一定比例执行cutmix和mixup增广

    cutmix = v2.CutMix(num_classes=len(classes))
    mixup = v2.MixUp(num_classes=len(classes))
    cutmix_or_mixup = v2.RandomChoice([cutmix, mixup])

    训练命令example

CUDA_VISIBLE_DEVICES=0 python3 -u train.py --backbone resnet101 --classnum 270 --workers 32 --lr=0.001 --epochs 30 --train_bs 160 --datadir your_dataset_directory

    简单的基于gradio的demo, app.py模型路径要简单调整一下

CUDA_VISIBLE_DEVICES=0 PORT=8000 python3 -u app.py

    如果训练过程,需要tensorboard显示中文,可以执行如下python指令,然后执行输出内容里的命令

python3 fixfont.py

/home/ubuntu/torch19/lib/python3.10/site-packages/matplotlib/mpl-data/matplotlibrc

/home/ubuntu/.cache/matplotlib

wget https://github.com/StellarCN/scp_zh/raw/master/fonts/SimHei.ttf

rm -f /home/ubuntu/.cache/matplotlib/*

cp ./SimHei.ttf /home/ubuntu/torch19/lib/python3.10/site-packages/matplotlib/mpl-data/fonts/ttf

你可能感兴趣的:(分类,数据挖掘,人工智能)