alexnet.py 1.6 KB

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