ذخیره و بارگذاری مدل در پایتورچ

در این بخش انواع روش های مورد استفاده، جهت ذخیره و بارگذاری مدل های PyTorch ارائه می شود.

چگونه مدل PyTorch را ذخیره و بارگذاری کنیم؟

وقتی صحبت از ذخیره و بارگذاری مدل ها می شود، سه عملکرد اصلی وجود دارد که باید با آنها آشنا شوید:

torch.save

torch.load

torch.nn.Module.load_state_dict

state_dict چیست؟

در PyTorch، پارامترهای قابل یادگیری (یعنی وزن ها و بایاس ها) torch.nn.Module در پارامترهای مدل موجود است (با ()model.parameters قابل دسترسی است). یک state_dict مانند یک شی در دیکشنری پایتون است که هر لایه را به تنسور پارامتر خود نگاشت می کند. فقط لایه‌هایی با پارامترهای قابل یادگیری (لایه‌های کانولوشن، لایه‌های خطی، و غیره) و registered buffers، ورودی‌هایی در state_dict مدل دارند. بهینه سازی (torch.optim) نیز دارای یک state_dict هستند. که حاوی اطلاعاتی در مورد وضعیت بهینه ساز و همچنین فراپارامترهای استفاده شده است.

از آنجایی که شی ها در state_dict دیکشنری پایتون هستند. می توان آنها را به راحتی ذخیره، به روز کرد، تغییر داد و بازیابی کرد و ماژولاریتی زیادی به مدل ها و بهینه سازهای PyTorch اضافه کرد.

مثال:

بیایید نگاهی به استفاده state_dict برای آموزش یک مدل ساده classifier بیاندازیم.

# Define model
class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# Initialize model
model = TheModelClass()

# Initialize optimizer
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# Print model's state_dict
print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

# Print optimizer's state_dict
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])

خروجی:

Model's state_dict:
conv1.weight     torch.Size([6, 3, 5, 5])
conv1.bias   torch.Size([6])
conv2.weight     torch.Size([16, 6, 5, 5])
conv2.bias   torch.Size([16])
fc1.weight   torch.Size([120, 400])
fc1.bias     torch.Size([120])
fc2.weight   torch.Size([84, 120])
fc2.bias     torch.Size([84])
fc3.weight   torch.Size([10, 84])
fc3.bias     torch.Size([10])

Optimizer's state_dict:
state    {}
param_groups     [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072, 4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]

ذخیره و بارگذاری مدل برای استنتاج:

ذخیره/بارگذاری state_dict (توصیه می‌شود)

ذخیره :

torch.save(model.state_dict(), PATH)

بارگذاری:

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

هنگام ذخیره یک مدل برای استنتاج، فقط لازم است پارامترهای مدل آموزش دیده ذخیره شود. ذخیره state_dict مدل با تابع ()torch.save بیشترین انعطاف را برای بازیابی مدل در آینده به شما می دهد، به همین دلیل است که روش توصیه شده برای ذخیره مدل ها است.

یک قرارداد رایج PyTorch ذخیره مدل ها با استفاده از پسوند  pt یا pth است.

قبل از اجرای استنتاج، باید () model.evalرا فراخوانی کنید تا لایه های dropout و batch normalizationرا روی حالت ارزیابی قرار دهید. عدم انجام این کار نتایج استنتاج متناقضی را به همراه خواهد داشت.

توجه داشته باشید که تابع ()load_state_dict یک شی دیکشنری می گیرد، نه مسیری به یک شی ذخیره شده. این بدان معناست که قبل از اینکه آن را به تابع ()load_state_dict منتقل کنید، باید state_dict ذخیره شده را deserialize کنید. به عنوان مثال، شما نمی توانید با استفاده model.load_state_dict(PATH) بارگذاری کنید.

ذخیره/بارگیری کل مدل:

ذخیره:

torch.save(model, PATH)

بارگذاری:

# Model class must be defined somewhere
model = torch.load(PATH)
model.eval()

این فرآیند ذخیره/بارگیری از بصری ترین نحو استفاده می کند و کمترین مقدار کد را شامل می شود. ذخیره یک مدل به این روش کل ماژول را با استفاده از ماژول pickle پایتون ذخیره می کند. ایراد این رویکرد این است که داده های سریال شده به کلاس های خاص و ساختار دایرکتوری دقیق مورد استفاده در هنگام ذخیره مدل محدود می شود. دلیل این امر این است که pickle خود کلاس مدل را ذخیره نمی کند. بلکه مسیری را برای فایل حاوی کلاس ذخیره می کند که در زمان بارگذاری استفاده می شود. به همین دلیل، کد شما هنگام استفاده در پروژه های دیگر یا پس از Refactor ها می تواند به روش های مختلف شکسته شود.

انتقال/بارگذاری مدل در قالب TorchScript:

یکی از راه‌های رایج برای استنتاج با یک مدل آموزش‌دیده، استفاده از TorchScript است، یک نمایش متوسط ​​از یک مدل PyTorch که می‌تواند در پایتون و همچنین در یک محیط با کارایی بالا مانند C++ اجرا شود. TorchScript در واقع فرمت مدل توصیه شده برای استنتاج و استقرار است.

با استفاده از قالب TorchScript، می‌توانید مدل export شده را بارگذاری کنید و استنتاج را بدون تعریف کلاس مدل اجرا کنید.

export:

model_scripted = torch.jit.script(model) # Export to TorchScript
model_scripted.save('model_scripted.pt') # Save

بارگذاری:

model = torch.jit.load('model_scripted.pt')
model.eval()

ذخیره و بارگذاری یک Checkpoint برای استنتاج و/یا از سرگیری آموزش :

ذخیره:

torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)

بارگذاری:

model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()
# - or -
model.train()

هنگام ذخیره یک checkpoint، برای استنباط یا از سرگیری آموزش، باید بیشتر از دستور state_dict استفاده کنید. ذخیره بهینه ساز state_dict نیز مهم است، زیرا شامل بافرها و پارامترهایی است که در اموزش مدل به روز می شوند. موارد دیگری که ممکن است بخواهید ذخیره کنید عبارتند از دوره ای که آن را متوقف کرده اید، آخرین lossآموزشی ثبت شده، لایه های خارجی torch.nn.Embedding و غیره است.

برای ذخیره چندین مؤلفه، آنها را در یک دیکشنری سازماندهی کنید و از ()torch.save برای سریال سازی دیکشنری لغت استفاده کنید. یک قرارداد رایج PyTorch ذخیره این checkpoint با استفاده از پسوند tar. است.

برای بارگذاری آیتم ها، ابتدا مدل و بهینه سازی را مقداردهی اولیه کنید، سپس دیکشنری را به صورت محلی با استفاده از (torch.load) بارگذاری کنید. از اینجا می‌توانید به سادگی با جستجو در دیکشنری همانطور که انتظار دارید به موارد ذخیره شده دسترسی پیدا کنید.

قبل از اجرای استنتاج، باید () model.eval را فراخوانی کنید تا لایه های dropout و نرمال سازی دسته ای را روی حالت ارزیابی قرار دهید. عدم انجام این کار نتایج استنتاج متناقضی را به همراه خواهد داشت. اگر می‌خواهید آموزش را از سر بگیرید، ()model.train را فراخوانی کنید تا مطمئن شوید که این لایه‌ها در حالت آموزش هستند.

ذخیره چندین مدل در یک فایل:

ذخیره:

torch.save({
            'modelA_state_dict': modelA.state_dict(),
            'modelB_state_dict': modelB.state_dict(),
            'optimizerA_state_dict': optimizerA.state_dict(),
            'optimizerB_state_dict': optimizerB.state_dict(),
            ...
            }, PATH)

بارگذاری:

modelA = TheModelAClass(*args, **kwargs)
modelB = TheModelBClass(*args, **kwargs)
optimizerA = TheOptimizerAClass(*args, **kwargs)
optimizerB = TheOptimizerBClass(*args, **kwargs)

checkpoint = torch.load(PATH)
modelA.load_state_dict(checkpoint['modelA_state_dict'])
modelB.load_state_dict(checkpoint['modelB_state_dict'])
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])

modelA.eval()
modelB.eval()
# - or -
modelA.train()
modelB.train()

هنگام ذخیره یک مدل متشکل از چندین torch.nn.Modules، مانند یک GAN، یک مدل sequence-to-sequence، یا مجموعه ای از مدل ها، از همان رویکردی که هنگام ذخیره یکcheckpoint استفاده می کنید، پیروی می کنید. به عبارت دیگر، یک دیکشنری از state_dict هر مدل و بهینه ساز مربوطه را ذخیره کنید. همانطور که قبلا ذکر شد، می توانید موارد دیگری را که ممکن است به شما در از سرگیری آموزش کمک کند، صرفاً با اضافه کردن آنها به دیکشنری ذخیره کنید.

Warmstarting Model با استفاده از پارامترهای مدل های مختلف:

ذخیره:

torch.save(modelA.state_dict(), PATH)

بارگذاری:

modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)

بارگذاری بخشی از یک مدل یا بارگذاری یک مدل جزئی سناریوهای رایج هنگام انتقال یادگیری یا آموزش یک مدل پیچیده جدید هستند. استفاده از پارامترهای آموزش دیده، حتی اگر تعداد کمی از آنها قابل استفاده باشند به روند آموزش کمک می کند. تا خیلی سریعتر از زمان اموزش از ابتدا همگرا شود.

چه از یک state_dict جزئی بارگیری کنید، که برخی از کلیدها را ندارد، یا یک state_dict را با کلیدهای بیشتری نسبت به مدلی که در آن بارگیری می‌کنید، بارگیری کنید. می‌توانید آرگومان strict را در تابع ()load_state_dict روی False تنظیم کنید تا کلیدهای غیر منطبق را نادیده بگیرید.

اگر می‌خواهید پارامترها را از یک لایه به لایه دیگر بارگیری کنید، اما برخی از کلیدها مطابقت ندارند، به سادگی نام کلیدهای پارامتر را در state_dict که بارگذاری می‌کنید تغییر دهید تا با کلیدهای مدلی که در آن بارگذاری می‌کنید مطابقت داشته باشد.

ذخیره و بارگذاری مدل در دیوایس ها:

ذخیره در GPU، بارگیری در CPU:

ذخیره:

torch.save(model.state_dict(), PATH)

بارگیری:

device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))

هنگام بارگذاری یک مدل بر روی CPU که با GPU آموزش دیده است. torch.device(‘cpu’) را به آرگومان map_location در تابع () torch.load ارسال کنید. در این مورد، ذخیره‌سازی‌های زیر تانسورها با استفاده از آرگومان map_location به صورت پویا به CPU نگاشت می‌شوند.

ذخیره در GPU، بارگیری در GPU:

ذخیره:

torch.save(model.state_dict(), PATH)

بارگذاری:

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)
# Make sure to call input = input.to(device) on any input tensors that you feed to the model

هنگام بارگذاری یک مدل روی GPU که  در GPU آموزش دیده ذخیره شده است. به سادگی مدل اولیه را با استفاده از model.to(torch.device(‘cuda’)) به یک مدل بهینه سازی شده CUDA تبدیل کنید. همچنین،از تابع my_tensor.to(device) در تمام ورودی های مدل برای آماده سازی داده های مدل استفاده کنید. توجه داشته باشید که فراخوانی my_tensor.to(device) یک کپی جدید از my_tensor را در GPU برمی گرداند. my_tensor را بازنویسی نمی کند. بنابراین، تانسورها را به صورت دستی بازنویسی کنید: my_tensor = my_tensor.to(torch.device(‘cuda’)).

ذخیره در CPU، بارگیری در GPU:

ذخیره:

torch.save(model.state_dict(), PATH)

بارگذاری:

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # Choose whatever GPU device number you want
model.to(device)
# Make sure to call input = input.to(device) on any input tensors that you feed to the model

هنگام بارگذاری یک مدل روی GPU که در CPU آموزش دیده و ذخیره شده . آرگومان map_location در تابع ()torch.load را روی cuda:device_id تنظیم کنید. این مدل را در یک GPU معین بارگذاری می کند. در مرحله بعد، حتما model.to(torch.device(‘cuda’)) را فراخوانی کنید تا تنسورهای پارامتر مدل را به تنسورهای CUDA تبدیل کنید.از تابع to(torch.device(‘cuda’)). در تمام ورودی های مدل برای آماده سازی داده ها برای مدل بهینه سازی شده CUDA استفاده کنید. توجه داشته باشید که فراخوانی  my_tensor.to(device)یک کپی جدید از my_tensor را در GPU برمی گرداند. my_tensor را بازنویسی نمی کند. بنابراین، به یاد داشته باشید که تانسورها را به صورت دستی بازنویسی کنید: my_tensor = my_tensor.to(torch.device(‘cuda’)).

ذخیره torch.nn.Data Models Parallel:

ذخیره:

.to(torch.device('cuda'))

بارگذاری:

# Load to whatever device you want

torch.nn.DataParallel استفاده از GPU موازی را امکان پذیر می کند. برای ذخیره یک مدل DataParallel به طور کلی، model.module.state_dict() را ذخیره کنید. به این ترتیب، شما این قابلیت را دارید که مدل را به هر شکلی که می خواهید در هر دیوایسی که می خواهید بارگذاری کنید.