PyTorch

import torch
import sol
import torchvision.models as models

''' Optimizing Model '''
py_model  = models.__dict__["alexnet"]()
input	  = torch.rand(32, 3, 224, 224)

# Use vdims=[True] if you plan to use changing batchsizes
sol_model = sol.optimize(py_model, vdims=[True])

''' Run training '''
sol_model.train()

# You cannot initialize the optimizer at this point. You need to wait until
# you have executed the model at least once, so SOL has compiled it.
optimizer = None
for batch in ...:
	input, target = ...
	output = sol_model(input)
	loss = loss_func(output, target)
	# After running the model once, you can safely initialize the optimizer
	if optimizer is None:
		optimizer = torch.optim.Adam(sol_model.parameters(), ...)
	optimizer.zero_grad()
	loss.backward()
	optimizer.step()
	...

''' Run validation '''
sol_model.eval()
with torch.no_grad():
	for batch in ...:
		input = ...
		output = sol_model(input)
		...

# Storing
sol_model = sol.optimize(pytorch_model, [...])
torch.save(sol_model.state_dict(), PATH)

# Loading
sol_model = sol.optimize(pytorch_model)
sol_model.load_state_dict(torch.load(PATH))

from torchvision import models

class Wrap(torch.nn.Module):
	def __init__(self, model):
		super().__init__()
		self.model = model

	def forward(self, x):
		out = self.model(x)
		if torch.jit.is_scripting():
			return (out[0], out[1]) if self.training else (out[0], None)
		return (out[0], out[1]) if self.training else (out, None)

model = Wrap(models.inception_v3())

# use only one output
model.training = False
sol_model = sol.optimize(model, ...)

# use two outputs
model.training = True
sol_model = sol.optimize(model, ...)

class ResNet50(pl.LightningModule):
	def __init__(self):
		super().__init__()
		self.model = torchvision.models.resnet50()

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

model = sol.optimize(ResNet50())