alexnet.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849
  1. from paddlex.cls import transforms
  2. import paddlex as pdx
  3. # 下载和解压蔬菜分类数据集
  4. veg_dataset = 'https://bj.bcebos.com/paddlex/datasets/vegetables_cls.tar.gz'
  5. pdx.utils.download_and_decompress(veg_dataset, path='./')
  6. # 定义训练和验证时的transforms
  7. train_transforms = transforms.Compose([
  8. transforms.RandomCrop(crop_size=224),
  9. transforms.RandomHorizontalFlip(),
  10. transforms.Normalize()
  11. ])
  12. eval_transforms = transforms.Compose([
  13. transforms.ResizeByShort(short_size=256),
  14. transforms.CenterCrop(crop_size=224),
  15. transforms.Normalize()
  16. ])
  17. # 定义训练和验证所用的数据集
  18. train_dataset = pdx.datasets.ImageNet(
  19. data_dir='vegetables_cls',
  20. file_list='vegetables_cls/train_list.txt',
  21. label_list='vegetables_cls/labels.txt',
  22. transforms=train_transforms,
  23. shuffle=True)
  24. eval_dataset = pdx.datasets.ImageNet(
  25. data_dir='vegetables_cls',
  26. file_list='vegetables_cls/val_list.txt',
  27. label_list='vegetables_cls/labels.txt',
  28. transforms=eval_transforms)
  29. # 初始化模型,并进行训练
  30. # 可使用VisualDL查看训练指标
  31. # VisualDL启动方式: visualdl --logdir output/mobilenetv2/vdl_log --port 8001
  32. # 浏览器打开 https://0.0.0.0:8001即可
  33. # 其中0.0.0.0为本机访问,如为远程服务, 改成相应机器IP
  34. model = pdx.cls.AlexNet(num_classes=len(train_dataset.labels))
  35. # AlexNet需要指定确定的input_shape
  36. model.fixed_input_shape = [224, 224]
  37. model.train(
  38. num_epochs=10,
  39. train_dataset=train_dataset,
  40. train_batch_size=32,
  41. eval_dataset=eval_dataset,
  42. lr_decay_epochs=[4, 6, 8],
  43. learning_rate=0.0025,
  44. save_dir='output/alexnet',
  45. use_vdl=True)