Created 3 years ago
# Imports
import torch
from torch.utils.data import Dataset, DataLoader
print(torch.__version__)
1.8.1
# NN APIs
import torch.autograd # Computation Graph
from torch.autograd import Variable # Variable node in computation graph
import torch.nn as nn # Neural Networks
import torch.nn.functional as F # Layers, activations, losses and more
import torch.optim as optim # optimizers: Adam, ...
# Distributed Training
import torch.distributed as dist # distributed communication
from multiprocessing import Process # memory sharing processes
# Torchvision
from torchvision import datasets, models, transforms # vision datasets, vision models, image transforms
import torchvision.transforms as transforms # composable transforms
# Torchscript and JIT
from torch.jit import script, trace # hybrid frontend decorator and tracing jit
torch.jit.trace(func, example_inputs)
# takes your module or function and an example data input, and traces the computational steps
# that the data encounters as it progresses through the model
@script # decorator used to indicate data-dependent control flow within the code being traced