分類

展開全部 | 收合全部

分類

展開全部 | 收合全部

【Pytorch】產生 single layer 的 Linear torch model, 並保存

內容目錄

前言

產生簡單的 torch model, 並保存

範例

產生一個最單純的 nn.Linear model, 並保存

import torch
import torch.nn as nn

class SingleLinearModel(nn.Module):
    def __init__(self):
        super(SingleLinearModel, self).__init__()
        self.linear = nn.Linear(in_features=10, out_features=5)

    def forward(self, x):
        return self.linear(x)

linear_model = SingleLinearModel()

dummy_input = torch.rand(1, 10)  # Batch size 1, input size 10

traced_script_module = torch.jit.trace(linear_model, dummy_input)
torch.jit.save(traced_script_module, 'model.pt')

print(" ------ model.pt save completed!!! ------")
Howard Weng
Howard Weng

我是 Howard Weng,很多人叫我嗡嗡。這個網站放了我的各種筆記。希望這些筆記也能順便幫助到有需要的人們!如果文章有幫助到你的話,歡迎幫我點讚哦!
另外,因為定位是「個人的隨手筆記」,有些文章內容「⚠️可能我理解有誤⚠️」或「?只寫到一半?」,如果有發現這樣的情況,歡迎在該文章的最下面留言提醒我!我會儘快修正或補上!感謝大家的建議與幫忙,讓網站能變得更好?

文章: 889