mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Update model prune loader to use pkgutil

Browse Source
This commit is contained in:
Ross Wightman 2023-02-06 17:45:16 -08:00
parent 0f2803de7a
commit 7a0bd095cb
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
timm/models/_prune.py
View File

@ -1,4 +1,5 @@
import os
import pkgutil
from copy import deepcopy
from torch import nn as nn
@ -108,6 +109,5 @@ def adapt_model_from_string(parent_module, model_string):
def adapt_model_from_file(parent_module, model_variant):
adapt_file = os.path.join(os.path.dirname(__file__), '_pruned', model_variant + '.txt')
with open(adapt_file, 'r') as f:
return adapt_model_from_string(parent_module, f.read().strip())
adapt_data = pkgutil.get_data(__name__, os.path.join('_pruned', model_variant + '.txt'))
return adapt_model_from_string(parent_module, adapt_data.decode('utf-8').strip())