Function main

CV/Pytorch_classification/ResNet/train.py:15–135

()

from model import resnet34

def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))

# 此处transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])中使用的是ImageNet数据集上图像RGB三通道的均值和方差
    
    # 此处定义的是train和val阶段分别使用的数据处理方法
    data_transform = {
        "train": transforms.Compose([transforms.RandomResizedCrop(224),
                                    transforms.RandomHorizontalFlip(),
                                    transforms.ToTensor(),
                                    transforms.Normalize([0.485, 0.456, 0.456], [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])])}


    # data_root = os.path.abspath(os.path.join(os.getcwd(), "../.."))  # 获取数据存储的根路径
    data_root = '/Users/WH/Desktop/pytorch_classification'
    print("current data path: {}".format(data_root))
    image_path = os.path.join(data_root, "flower_data")
    assert os.path.exists(image_path), "{} path does not exists.".format(image_path)
    train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"),
                                         transform=data_transform["train"])
    train_num = len(train_dataset)

    # {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4}

    flower_list = train_dataset.class_to_idx
    cla_dict = dict((val, key) for key, val in flower_list.items())
    # write dict into json file
    json_str = json.dumps(cla_dict, indent=4)
    with open('class_indices.json', 'w') as json_file:
        json_file.write(json_str)

    batch_size = 16
    nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # num of workes
    print('Using {} dataloader workers every process'.format(nw))

    train_loader = torch.utils.data.DataLoader(train_dataset,
                                              batch_size=batch_size, shuffle=True,
                                              num_workers=nw)
    
    validate_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val"),
                                            transform=data_transform["val"])
    val_num = len(validate_dataset)
    validate_loader = torch.utils.data.DataLoader(validate_dataset,
                                                  batch_size=batch_size, shuffle=False,
                                                  num_workers=nw)
    
    print("using {} images for training, {} images for validation.".format(train_num,
                                                                           val_num))


    net = resnet34() # 模型实例化   

    # load pretrain weights