万州网站制作沈阳网站关键词优化公司
最近需要训练一个有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
- ...
- ...
- class1
自定义一个Dataset，实现如下
class ImageFolderEx(Dataset):def __init__(self, image_dir, image_files, image_labels, classnum=1000, transform=None):self.image_dir = image_dirself.image_files = image_filesself.image_labels = image_labelsself.classnum = classnumself.transform = transformdef __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.classnumonehot[label] = 1return (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