pytorch模型转换为android支持的模型

1、模型转换（从gpu到cpu）

如过在训练的时候采用的是gpu训练的模型，我们需要保存为cpu格式的模型，如下代码所示：

def from_gpu_model_to_cpu_model(read_gpu_pth,write_cpu_pth):
    device = torch.device('cpu')
    model=torch.load(read_gpu_pth, map_location = device)
    torch.save(model,write_cpu_pth)
    print('from gpu model to cpu model success!')

2、模型转换（从windows到android）

def from_win_pytorch_to_android(read_win_model_path, write_android_model_path):
    model = torch.load(read_win_model_path)
    model.eval()
    x = torch.rand(1, 3, 96, 96)
    script_model = torch.jit.trace(model,x)
    script_model.save(write_android_model_path)
    print('android model save success!')