russian_art_2024

1
"""Script model to ONNX format"""
2

3
import torch
4

5
from src.initial_model_utils import init_model
6

7
MODEL_WEIGHTS = "./data/weights/resnet50_tl_68.pt"
8

9
device = torch.device("cpu")
10
model = init_model(device, num_classes=35)
11
model.load_state_dict(torch.load(MODEL_WEIGHTS, map_location=device))
12
model.eval()
13

14
# TODO
15
# Дописать конвертацию в ONNX
16