Метод code

Метод code класса ScriptModule возвращает строку, содержащую исходный код модуля, который был скомпилирован с помощью torch.jit.script. Этот метод полезен для отладки и понимания того, как PyTorch JIT преобразует Python-код в промежуточное представление.

Синтаксис

script_module.code

Метод не принимает параметров и возвращает строку с исходным кодом.

Пример

Давайте создадим простой модуль и скомпилируем его с помощью torch.jit.script, а затем получим его исходный код:

import torch

class MyModule(torch.nn.Module):
	def forward(self, x):
		return x + 1

scripted = torch.jit.script(MyModule())
print(scripted.code)

Результат выполнения кода:

"def forward(self,
	x: Tensor) -> Tensor:
  return x + 1"

Пример

Рассмотрим более сложный модуль с несколькими операциями и условным выражением:

import torch

class ComplexModule(torch.nn.Module):
	def forward(self, x, y):
		if x.sum() > y.sum():
			return x * 2
		else:
			return y * 3

scripted = torch.jit.script(ComplexModule())
print(scripted.code)

Результат выполнения кода:

"def forward(self,
	x: Tensor,
	y: Tensor) -> Tensor:
  if torch.sum(x) > torch.sum(y):
	return x * 2
  else:
	return y * 3"

Пример

Метод code также работает с модулями, содержащими вложенные подмодули. Полученный код показывает все операции, включая вызовы вложенных модулей:

import torch

class SubModule(torch.nn.Module):
	def forward(self, x):
		return x * 2

class MainModule(torch.nn.Module):
	def __init__(self):
		super().__init__()
		self.sub = SubModule()
	
	def forward(self, x):
		return self.sub(x) + 1

scripted = torch.jit.script(MainModule())
print(scripted.code)

Результат выполнения кода:

"def forward(self,
	x: Tensor) -> Tensor:
  sub_0 = self.sub
  _0 = (sub_0).forward(x, )
  return _0 + 1"

Смотрите также

класс ScriptModule,
который представляет скомпилированный модуль в PyTorch JIT
метод save,
который сохраняет скомпилированный модуль на диск
метод graph,
который возвращает вычислительный граф модуля
метод code,
который возвращает исходный код скомпилированного модуля