资讯 社区 文档
技术能力
语音技术
文字识别
人脸与人体
图像技术
语言与知识
视频技术

Pytorch 1.7.1代码规范

Pytorch 1.7.1代码规范

基于Pytorch 1.7.1框架的MNIST图像分类,训练数据集pytorch_train_data.zip点击这里下载。
如下所示是其超参搜索任务中一个超参数组合的训练代码,代码会通过argparse模块接受在平台中填写的信息,请保持一致。
特别注意,示例采用的是进化算法进行超参搜索,每个试验在训练时会继承之前试验的权重,resume_checkpoint_path是权重的保存路径,由搜索算法自身提供,与job_id及trial_id一样,只需要在argparse中提供对应参数即可。

pytorch1.7.1_autosearch.py示例代码

# -*- coding:utf-8 -*-
""" pytorch train demo """
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data as data
from torchvision import transforms
import codecs
import errno
import gzip
import numpy as np
import os
import time
from PIL import Image
from rudder_autosearch.sdk.amaas_tools import AMaasTools

def parse_arg():
    """parse arguments"""
    parser = argparse.ArgumentParser(description='PyTorch1.7.1 MNIST Example')
    parser.add_argument('--train_dir', type=str, default='./train_data',
                        help='input data dir for training (default: ./train_data)')
    parser.add_argument('--test_dir', type=str, default='./test_data',
                        help='input data dir for test (default: ./test_data)')
    parser.add_argument('--output_dir', type=str, default='./output',
                        help='output dir for auto_search job (default: ./output)')
    parser.add_argument('--job_id', type=str, default="job-1234",
                        help='auto_search job id (default: "job-1234")')
    parser.add_argument('--trial_id', type=str, default="0-0",
                        help='auto_search id of a single trial (default: "0-0")')
    parser.add_argument('--metric', type=str, default="acc",
                        help='evaluation metric of the model')
    parser.add_argument('--data_sampling_scale', type=float, default=1.0,
                        help='sampling ratio of the data (default: 1.0)')
    parser.add_argument('--batch_size', type=int, default=64,
                        help='number of images input in an iteration (default: 64)')
    parser.add_argument('--lr', type=float, default=0.01,
                        help='learning rate (default: 0.01)')
    parser.add_argument('--momentum', type=float, default=0.5,
                        help='SGD momentum (default: 0.5)')
    parser.add_argument('--no_cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--log_interval', type=int, default=10,
                        help='how many batches to wait before logging training status')
    parser.add_argument('--perturb_interval', type=int, default=10,
                        help='number of epochs to train (default: 10)')
    parser.add_argument('--resume_checkpoint_path', type=str, default="",
                        help='inherit the initial weight of the previous trial')
    args = parser.parse_args()
    args.output_dir = os.path.join(args.output_dir, args.job_id, args.trial_id)
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)
    args.cuda = not args.no_cuda and torch.cuda.is_available()


    print("job_id: {}, trial_id: {}".format(args.job_id, args.trial_id))
    return args

# 定义MNIST数据集的dataset
class MNIST(data.Dataset):
    """
    MNIST dataset
    """
    training_file = 'training.pt'
    test_file = 'test.pt'
    classes = ['0 - zero', '1 - one', '2 - two', '3 - three', '4 - four',
               '5 - five', '6 - six', '7 - seven', '8 - eight', '9 - nine']

    def __init__(self, root, train=True, transform=None, target_transform=None, data_sampling_scale=1):
        self.root = os.path.expanduser(root)
        self.transform = transform
        self.target_transform = target_transform
        self.train = train  # training set or test set
        self.data_sampling_scale = data_sampling_scale
        self.preprocess(root, train, False)
        if self.train:
            data_file = self.training_file
        else:
            data_file = self.test_file
        self.data, self.targets = torch.load(os.path.join(self.processed_folder, data_file))

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        img, target = self.data[index], int(self.targets[index])
        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(img.numpy(), mode='L')
        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

    def __len__(self):
        return len(self.data)

    @property
    def raw_folder(self):
        """
        raw folder
        """
        return os.path.join('/tmp', 'raw')

    @property
    def processed_folder(self):
        """
        processed folder
        """
        return os.path.join('/tmp', 'processed')

    # data preprocessing
    def preprocess(self, train_dir, train, remove_finished=False):
        """
        preprocess
        """
        makedir_exist_ok(self.raw_folder)
        makedir_exist_ok(self.processed_folder)
        train_list = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz']
        test_list = ['t10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
        zip_list = train_list if train else test_list
        for zip_file in zip_list:
            print('Extracting {}'.format(zip_file))
            zip_file_path = os.path.join(train_dir, zip_file)
            raw_folder_path = os.path.join(self.raw_folder, zip_file)
            with open(raw_folder_path.replace('.gz', ''), 'wb') as out_f, gzip.GzipFile(zip_file_path) as zip_f:
                out_f.write(zip_f.read())
            if remove_finished:
                os.unlink(zip_file_path)
        if train:
            x_train = read_image_file(os.path.join(self.raw_folder, 'train-images-idx3-ubyte'))
            y_train = read_label_file(os.path.join(self.raw_folder, 'train-labels-idx1-ubyte'))
            np.random.seed(0)
            sample_data_num = int(self.data_sampling_scale * len(x_train))
            idx = np.arange(len(x_train))
            np.random.shuffle(idx)
            x_train, y_train = x_train[0:sample_data_num], y_train[0:sample_data_num]
            training_set = (x_train, y_train)
            with open(os.path.join(self.processed_folder, self.training_file), 'wb') as f:
                torch.save(training_set, f)
        else:
            test_set = (
                read_image_file(os.path.join(self.raw_folder, 't10k-images-idx3-ubyte')),
                read_label_file(os.path.join(self.raw_folder, 't10k-labels-idx1-ubyte'))
            )
            with open(os.path.join(self.processed_folder, self.test_file), 'wb') as f:
                torch.save(test_set, f)

def get_int(b):
    """
    get int
    """
    return int(codecs.encode(b, 'hex'), 16)

def read_label_file(path):
    """
    read label file
    """
    with open(path, 'rb') as f:
        data = f.read()
        assert get_int(data[:4]) == 2049
        length = get_int(data[4:8])
        parsed = np.frombuffer(data, dtype=np.uint8, offset=8)
        return torch.from_numpy(parsed).view(length).long()

def read_image_file(path):
    """
    read image file
    """
    with open(path, 'rb') as f:
        data = f.read()
        assert get_int(data[:4]) == 2051
        length = get_int(data[4:8])
        num_rows = get_int(data[8:12])
        num_cols = get_int(data[12:16])
        parsed = np.frombuffer(data, dtype=np.uint8, offset=16)
        return torch.from_numpy(parsed).view(length, num_rows, num_cols)

def makedir_exist_ok(dirpath):
    """
    Python2 support for os.makedirs(.., exist_ok=True)
    """
    try:
        os.makedirs(dirpath)
    except OSError as e:
        if e.errno == errno.EEXIST:
            pass
        else:
            raise

def load_data(args):
    """load_data"""
    # 若无测试集,训练集做验证集
    if not os.path.exists(args.test_dir) or not os.listdir(args.test_dir):
        args.test_dir = args.train_dir
    # 将数据进行转化,从PIL.Image/numpy.ndarray的数据进转化为torch.FloadTensor
    trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    train_set = MNIST(root=args.train_dir, train=True, transform=trans, data_sampling_scale=args.data_sampling_scale)
    test_set = MNIST(root=args.test_dir, train=False, transform=trans)
    # 定义data reader
    train_loader = torch.utils.data.DataLoader(
        dataset=train_set,
        batch_size=args.batch_size,
        shuffle=True)
    test_loader = torch.utils.data.DataLoader(
        dataset=test_set,
        batch_size=args.batch_size,
        shuffle=False)
    return train_loader, test_loader

# 定义网络模型
class Net(nn.Module):
    """
    Net
    """
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        """
        forward
        """
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x)

def load_state_dict(model, resume_checkpoint_path):
    """load_state_dict"""
    if resume_checkpoint_path:
        model.load_state_dict(torch.load(resume_checkpoint_path))

def run_train(model, args, train_loader):
    """run_train"""
    if args.cuda:
        # Move model to GPU.
        model.cuda()
    # 选择优化器
    optimizer = optim.SGD(model.parameters(), lr=args.lr,
                          momentum=args.momentum)
    for epoch in range(1, args.perturb_interval + 1):
        train(model, args, train_loader, optimizer, epoch)

def train(model, args, train_loader, optimizer, epoch):
    """
    train
    """
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        if args.cuda:
            data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()
        output = model(data)  # 获取预测值
        loss = F.nll_loss(output, target)  # 计算loss
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx, len(train_loader),
                100. * batch_idx / len(train_loader), loss.item()))

def evaluate(model, args, test_loader):
    """evaluate"""
    model.eval()
    test_loss = 0.
    test_accuracy = 0.
    for data, target in test_loader:
        if args.cuda:
            data, target = data.cuda(), target.cuda()
        output = model(data)
        # sum up batch loss
        test_loss += F.nll_loss(output, target, size_average=False).item()
        # get the index of the max log-probability
        pred = output.data.max(1, keepdim=True)[1]
        test_accuracy += pred.eq(target.data.view_as(pred)).cpu().float().sum()
    test_loss /= len(test_loader) * args.batch_size
    test_accuracy /= len(test_loader) * args.batch_size
    print('\nTest set: Average loss: {:.4f}, Accuracy: {:.2f}%\n'.format(
        test_loss, 100. * test_accuracy))
    return float(test_accuracy)

def save(model, output_dir):
    """
    save
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    # 保存模型
    torch.save(model.state_dict(), os.path.join(output_dir, 'model.pkl'))

def report_final(args, metric):
    """report_final_result"""
    # 结果上报sdk
    amaas_tools = AMaasTools(args.job_id, args.trial_id)
    metric_dict = {args.metric: metric}
    checkpoint_path = os.path.join(args.output_dir, 'model.pkl')
    for i in range(3):
        flag, ret_msg = amaas_tools.report_final_result(metric=metric_dict,
                                                        export_model_path=args.output_dir,
                                                        checkpoint_path=checkpoint_path)
        print("End Report, metric:{}, ret_msg:{}".format(metric, ret_msg))
        if flag:
            break
        time.sleep(1)
    assert flag, "Report final result to manager failed! Please check whether manager'address or manager'status " \
                 "is ok! "

def main():
    """main"""
    # 获取参数
    args = parse_arg()
    # 加载数据集
    train_loader, test_loader = load_data(args)
    # 模型定义
    model = Net()
    # 继承之前实验的模型参数
    load_state_dict(model, args.resume_checkpoint_path)
    # 模型训练
    run_train(model, args, train_loader)
    # 模型保存
    save(model, args.output_dir)
    # 模型评估
    acc = evaluate(model, args, test_loader)
    # 上报结果
    report_final(args, metric=acc)

if __name__ == '__main__':
    main()

示例代码对应的yaml配置如下,请保持格式一致

pbt_search_demo.yml示例内容

#搜索算法参数
search_strategy:
  algo: PBT_SEARCH #搜索策略:进化算法
  params:
    population_num: 8 #种群个体数量 | [1,10] int类型
    round: 10 #迭代轮数  |[5,50] int类型
    perturb_interval: 10 # 扰动间隔  | [1,20] int类型
    quantile_frac: 0.5 #扰动比例  | (0,0.5] float类型
    explore_prob: 0.25 #扰动概率 | (0,0.5] float类型

#单次训练时数据的采样比例,单位%
data_sampling_scale: 100 #|(0,100] int类型

#评价指标参数
metrics:
  name: acc #评价指标 | 任意字符串 str类型
  goal: MAXIMIZE #最大值/最小值 | str类型   MAXIMIZE or MINIMIZE   必须为这两个之一(也即支持大写)
  expected_value: 100 #早停标准值,评价指标超过该值则结束整个超参搜索,单位% |无限制 int类型

#搜索参数空间
search_space:
  batch_size:
    htype: choice
    value: [64, 128, 256, 512]
  lr:
    htype: loguniform
    value: [0.0001, 0.1]
  momentum:
    htype: uniform
    value: [0.1, 0.9]

pytorch_predict.py示例代码

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@license: Copyright (c) 2019 Baidu.com, Inc. All Rights Reserved.
@desc: 图像预测算法示例
"""
import logging
import torch
import torch.nn as nn
import torch.nn.functional as F
import base64
import json
from PIL import Image
from io import BytesIO
from torchvision import datasets, models, transforms
MODEL_FILE_NAME = 'model.pkl'  # 模型文件名称
def get_image_transform():
    """获取图片处理的transform
    Args:
        data_type: string, type of data(train/test)
    Returns:
        torchvision.transforms.Compose
    """
    trans = transforms.Compose([transforms.Resize((28, 28)),
                                transforms.ToTensor(),
                                transforms.Normalize((0.5,), (1.0,))])
    return trans
def model_fn(model_dir):
    """模型加载
    Args:
        model_dir: 模型路径,该目录存储的文件为在自动搜索作业中选择的输出路径下产出的文件
    Returns:
        加载好的模型对象
    """
    class Net(nn.Module):
        """Net"""
        def __init__(self):
            super(Net, self).__init__()
            self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
            self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
            self.conv2_drop = nn.Dropout2d()
            self.fc1 = nn.Linear(320, 50)
            self.fc2 = nn.Linear(50, 10)
        def forward(self, x):
            """
            forward
            """
            x = F.relu(F.max_pool2d(self.conv1(x), 2))
            x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
            x = x.view(-1, 320)
            x = F.relu(self.fc1(x))
            x = F.dropout(x, training=self.training)
            x = self.fc2(x)
            return F.log_softmax(x)
    model = Net()
    meta_info_path = "%s/%s" % (model_dir, MODEL_FILE_NAME)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.load_state_dict(torch.load(meta_info_path, map_location=device))
    model.to(device)
    logging.info("device type: " + str(device))
    return model
def input_fn(request):
    """对输入进行格式化,处理为预测需要的输入格式
    Args:
        request: api请求的json
    Returns:
        预测需要的输入数据,一般为tensor
    """
    instances = request['instances']
    transform_composes = get_image_transform()
    arr_tensor_data = []
    for instance in instances:
        decoded_data = base64.b64decode(instance['data'].encode("utf8"))
        byte_stream = BytesIO(decoded_data)
        roiImg = Image.open(byte_stream)
        target_data = transform_composes(roiImg)
        arr_tensor_data.append(target_data)
    tensor_data = torch.stack(arr_tensor_data, dim=0)
    return tensor_data
def output_fn(predict_result):
    """进行输出格式化
    Args:
        predict_result: 预测结果
    Returns:
        格式化后的预测结果,需能够json序列化以便接口返回
    """
    js_str = None
    if type(predict_result) == torch.Tensor:
        list_prediction = predict_result.detach().cpu().numpy().tolist()
        js_str = json.dumps(list_prediction)
    return js_str
上一篇
Sklearn 0.23.2代码规范
下一篇
Tensorflow2.3.0代码规范