最近看了常用的图像分类网络,分别使用迁移学习和直接构建模型自己训练。找一个模型把使用过程写下来吧。
1、加载数据集并预处理
batch_size = 50
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# data pre-treatment
data_transform = {
"train": transforms.Compose([transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
"val": transforms.Compose([transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]