以下是一个使用AWS Lambda函数运行一个小的PyTorch模型的示例代码:
lambda_function.py的Python文件,并将以下代码复制到该文件中:import torch
# 定义一个简单的PyTorch模型
class SimpleModel(torch.nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc = torch.nn.Linear(1, 1)
def forward(self, x):
return self.fc(x)
def lambda_handler(event, context):
# 创建模型实例
model = SimpleModel()
# 加载训练好的模型权重
model.load_state_dict(torch.load('model_weights.pth'))
# 输入数据
input_data = torch.tensor([[2.0]], dtype=torch.float32)
# 使用模型进行推理
output_data = model(input_data)
# 返回推理结果
return {
'statusCode': 200,
'body': output_data.item()
}
然后,创建一个名为model_weights.pth的文件,并将训练好的模型权重保存到该文件中。请确保模型权重文件与lambda_function.py在同一个目录下。
接下来,使用AWS CLI命令将代码部署到AWS Lambda函数中。打开终端并运行以下命令:
aws lambda create-function --function-name my-lambda-function --runtime python3.8 --zip-file fileb://lambda_function.zip --handler lambda_function.lambda_handler --role your-iam-role-arn
请将my-lambda-function替换为您的函数名称,your-iam-role-arn替换为您的IAM角色的ARN。
aws lambda invoke --function-name my-lambda-function output.json
将my-lambda-function替换为您的函数名称。
以上示例代码演示了如何在AWS Lambda函数中运行一个简单的PyTorch模型。您可以根据自己的需求修改代码,并根据您的模型进行相应的推理操作。