def _test(test_case):
if os.getenv("ONEFLOW_TEST_CPU_ONLY"):
device = flow.device("cpu")
else:
device = flow.device("cuda")
net = Net()
net.to(device)
optimizer = optim.SGD(net.parameters(), lr=0.002, momentum=0.9)
criterion = nn.CrossEntropyLoss()
criterion.to(device)
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
train_epoch = 1
batch_size = 4
num_workers = 0
data_dir = os.path.join(os.getenv("ONEFLOW_TEST_CACHE_DIR", "./data-test"),
"cifar10")
train_iter, test_iter = load_data_cifar10(
batch_size=batch_size,
data_dir=data_dir,
download=True,
transform=transform,
source_url=
"https://oneflow-public.oss-cn-beijing.aliyuncs.com/datasets/cifar/cifar-10-python.tar.gz",
num_workers=num_workers,
)
final_loss = 0
for epoch in range(1, train_epoch +
1): # loop over the dataset multiple times
running_loss = 0.0
for i, data in enumerate(train_iter, 1):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
inputs = inputs.to(dtype=flow.float32, device=device)
labels = labels.to(dtype=flow.int64, device=device)
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % 200 == 0: # print every 200 mini-batches
final_loss = running_loss / 200
print("epoch: %d step: %5d loss: %.3f " %
(epoch, i, final_loss))
running_loss = 0.0
break
print("final loss : ", final_loss)
# -----------------
# Classic cat example!
from PIL import Image
img_url = "https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true"
img_path = download_testdata(img_url, "cat.png", module="data")
img = Image.open(img_path).resize((224, 224))
# Preprocess the image and convert to tensor
from flowvision import transforms
my_preprocess = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224,
0.225]),
])
img = my_preprocess(img)
img = np.expand_dims(img.numpy(), 0)
######################################################################
# Import the graph to Relay
# -------------------------
# Convert OneFlow graph to Relay graph. The input name can be arbitrary.
class Graph(flow.nn.Graph):
def __init__(self, module):
super().__init__()
self.m = module
def build(self, x):