Python caffe2.python.core.DeviceScope() Examples
The following are 30
code examples of caffe2.python.core.DeviceScope().
You can vote up the ones you like or vote down the ones you don't like,
and go to the original project or source file by following the links above each example.
You may also want to check out all available functions/classes of the module
caffe2.python.core
, or try the search function
.
.
Example #1
| Source File: test_spatial_narrow_as_op.py From CBNet with Apache License 2.0 | 6 votes |
|
def _run_test(self, A, B, check_grad=False):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
op = core.CreateOperator('SpatialNarrowAs', ['A', 'B'], ['C'])
workspace.FeedBlob('A', A)
workspace.FeedBlob('B', B)
workspace.RunOperatorOnce(op)
C = workspace.FetchBlob('C')
if check_grad:
gc = gradient_checker.GradientChecker(
stepsize=0.005,
threshold=0.005,
device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
)
res, grad, grad_estimated = gc.CheckSimple(op, [A, B], 0, [0])
self.assertTrue(res, 'Grad check failed')
dims = C.shape
C_ref = A[:dims[0], :dims[1], :dims[2], :dims[3]]
np.testing.assert_allclose(C, C_ref, rtol=1e-5, atol=1e-08)
Example #2
| Source File: test_batch_permutation_op.py From KL-Loss with Apache License 2.0 | 6 votes |
|
def _run_op_test(self, X, I, check_grad=False):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
op = core.CreateOperator('BatchPermutation', ['X', 'I'], ['Y'])
workspace.FeedBlob('X', X)
workspace.FeedBlob('I', I)
workspace.RunOperatorOnce(op)
Y = workspace.FetchBlob('Y')
if check_grad:
gc = gradient_checker.GradientChecker(
stepsize=0.1,
threshold=0.001,
device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
)
res, grad, grad_estimated = gc.CheckSimple(op, [X, I], 0, [0])
self.assertTrue(res, 'Grad check failed')
Y_ref = X[I]
np.testing.assert_allclose(Y, Y_ref, rtol=1e-5, atol=1e-08)
Example #3
| Source File: test_loader.py From KL-Loss with Apache License 2.0 | 6 votes |
|
def get_net(data_loader, name):
logger = logging.getLogger(__name__)
blob_names = data_loader.get_output_names()
net = core.Net(name)
net.type = 'dag'
for gpu_id in range(cfg.NUM_GPUS):
with core.NameScope('gpu_{}'.format(gpu_id)):
with core.DeviceScope(muji.OnGPU(gpu_id)):
for blob_name in blob_names:
blob = core.ScopedName(blob_name)
workspace.CreateBlob(blob)
net.DequeueBlobs(
data_loader._blobs_queue_name, blob_names)
logger.info("Protobuf:\n" + str(net.Proto()))
return net
Example #4
| Source File: test_spatial_narrow_as_op.py From KL-Loss with Apache License 2.0 | 6 votes |
|
def _run_test(self, A, B, check_grad=False):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
op = core.CreateOperator('SpatialNarrowAs', ['A', 'B'], ['C'])
workspace.FeedBlob('A', A)
workspace.FeedBlob('B', B)
workspace.RunOperatorOnce(op)
C = workspace.FetchBlob('C')
if check_grad:
gc = gradient_checker.GradientChecker(
stepsize=0.005,
threshold=0.005,
device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
)
res, grad, grad_estimated = gc.CheckSimple(op, [A, B], 0, [0])
self.assertTrue(res, 'Grad check failed')
dims = C.shape
C_ref = A[:dims[0], :dims[1], :dims[2], :dims[3]]
np.testing.assert_allclose(C, C_ref, rtol=1e-5, atol=1e-08)
Example #5
| Source File: bn_helper.py From video-long-term-feature-banks with Apache License 2.0 | 6 votes |
|
def _update_bn_stats_gpu(self):
"""
Copy to GPU.
Note: the actual blobs used at test time are "rm" and "riv"
"""
num_gpus = cfg.NUM_GPUS
root_gpu_id = cfg.ROOT_GPU_ID
for i in range(root_gpu_id, root_gpu_id + num_gpus):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, i)):
for bn_layer in self._bn_layers:
workspace.FeedBlob(
'gpu_{}/'.format(i) + bn_layer + '_bn_rm',
np.array(self._meanX_dict[bn_layer], dtype=np.float32),
)
"""
Note: riv is acutally running var (not running inv var)!!!!
"""
workspace.FeedBlob(
'gpu_{}/'.format(i) + bn_layer + '_bn_riv',
np.array(self._var_dict[bn_layer], dtype=np.float32),
)
Example #6
| Source File: test_spatial_narrow_as_op.py From Clustered-Object-Detection-in-Aerial-Image with Apache License 2.0 | 6 votes |
|
def _run_test(self, A, B, check_grad=False):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
op = core.CreateOperator('SpatialNarrowAs', ['A', 'B'], ['C'])
workspace.FeedBlob('A', A)
workspace.FeedBlob('B', B)
workspace.RunOperatorOnce(op)
C = workspace.FetchBlob('C')
if check_grad:
gc = gradient_checker.GradientChecker(
stepsize=0.005,
threshold=0.005,
device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
)
res, grad, grad_estimated = gc.CheckSimple(op, [A, B], 0, [0])
self.assertTrue(res, 'Grad check failed')
dims = C.shape
C_ref = A[:dims[0], :dims[1], :dims[2], :dims[3]]
np.testing.assert_allclose(C, C_ref, rtol=1e-5, atol=1e-08)
Example #7
| Source File: test_batch_permutation_op.py From Clustered-Object-Detection-in-Aerial-Image with Apache License 2.0 | 6 votes |
|
def _run_op_test(self, X, I, check_grad=False):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
op = core.CreateOperator('BatchPermutation', ['X', 'I'], ['Y'])
workspace.FeedBlob('X', X)
workspace.FeedBlob('I', I)
workspace.RunOperatorOnce(op)
Y = workspace.FetchBlob('Y')
if check_grad:
gc = gradient_checker.GradientChecker(
stepsize=0.1,
threshold=0.001,
device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
)
res, grad, grad_estimated = gc.CheckSimple(op, [X, I], 0, [0])
self.assertTrue(res, 'Grad check failed')
Y_ref = X[I]
np.testing.assert_allclose(Y, Y_ref, rtol=1e-5, atol=1e-08)
Example #8
| Source File: test_loader.py From Clustered-Object-Detection-in-Aerial-Image with Apache License 2.0 | 6 votes |
|
def get_net(data_loader, name):
logger = logging.getLogger(__name__)
blob_names = data_loader.get_output_names()
net = core.Net(name)
net.type = 'dag'
for gpu_id in range(cfg.NUM_GPUS):
with core.NameScope('gpu_{}'.format(gpu_id)):
with core.DeviceScope(muji.OnGPU(gpu_id)):
for blob_name in blob_names:
blob = core.ScopedName(blob_name)
workspace.CreateBlob(blob)
net.DequeueBlobs(
data_loader._blobs_queue_name, blob_names)
logger.info("Protobuf:\n" + str(net.Proto()))
return net
Example #9
| Source File: test_spatial_narrow_as_op.py From seg_every_thing with Apache License 2.0 | 6 votes |
|
def _run_test(self, A, B, check_grad=False):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
op = core.CreateOperator('SpatialNarrowAs', ['A', 'B'], ['C'])
workspace.FeedBlob('A', A)
workspace.FeedBlob('B', B)
workspace.RunOperatorOnce(op)
C = workspace.FetchBlob('C')
if check_grad:
gc = gradient_checker.GradientChecker(
stepsize=0.005,
threshold=0.005,
device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
)
res, grad, grad_estimated = gc.CheckSimple(op, [A, B], 0, [0])
self.assertTrue(res, 'Grad check failed')
dims = C.shape
C_ref = A[:dims[0], :dims[1], :dims[2], :dims[3]]
np.testing.assert_allclose(C, C_ref, rtol=1e-5, atol=1e-08)
Example #10
| Source File: test_batch_permutation_op.py From seg_every_thing with Apache License 2.0 | 6 votes |
|
def _run_op_test(self, X, I, check_grad=False):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
op = core.CreateOperator('BatchPermutation', ['X', 'I'], ['Y'])
workspace.FeedBlob('X', X)
workspace.FeedBlob('I', I)
workspace.RunOperatorOnce(op)
Y = workspace.FetchBlob('Y')
if check_grad:
gc = gradient_checker.GradientChecker(
stepsize=0.1,
threshold=0.001,
device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
)
res, grad, grad_estimated = gc.CheckSimple(op, [X, I], 0, [0])
self.assertTrue(res, 'Grad check failed')
Y_ref = X[I]
np.testing.assert_allclose(Y, Y_ref, rtol=1e-5, atol=1e-08)
Example #11
| Source File: test_loader.py From seg_every_thing with Apache License 2.0 | 6 votes |
|
def get_net(data_loader, name):
logger = logging.getLogger(__name__)
blob_names = data_loader.get_output_names()
net = core.Net(name)
net.type = 'dag'
for gpu_id in range(cfg.NUM_GPUS):
with core.NameScope('gpu_{}'.format(gpu_id)):
with core.DeviceScope(muji.OnGPU(gpu_id)):
for blob_name in blob_names:
blob = core.ScopedName(blob_name)
workspace.CreateBlob(blob)
net.DequeueBlobs(
data_loader._blobs_queue_name, blob_names)
logger.info("Protobuf:\n" + str(net.Proto()))
return net
Example #12
| Source File: test_spatial_narrow_as_op.py From masktextspotter.caffe2 with Apache License 2.0 | 6 votes |
|
def _run_test(self, A, B, check_grad=False):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
op = core.CreateOperator('SpatialNarrowAs', ['A', 'B'], ['C'])
workspace.FeedBlob('A', A)
workspace.FeedBlob('B', B)
workspace.RunOperatorOnce(op)
C = workspace.FetchBlob('C')
if check_grad:
gc = gradient_checker.GradientChecker(
stepsize=0.005,
threshold=0.005,
device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
)
res, grad, grad_estimated = gc.CheckSimple(op, [A, B], 0, [0])
self.assertTrue(res, 'Grad check failed')
dims = C.shape
C_ref = A[:dims[0], :dims[1], :dims[2], :dims[3]]
np.testing.assert_allclose(C, C_ref, rtol=1e-5, atol=1e-08)
Example #13
| Source File: test_loader.py From masktextspotter.caffe2 with Apache License 2.0 | 6 votes |
|
def get_net(data_loader, name):
logger = logging.getLogger(__name__)
blob_names = data_loader.get_output_names()
net = core.Net(name)
net.type = 'dag'
for gpu_id in range(cfg.NUM_GPUS):
with core.NameScope('gpu_{}'.format(gpu_id)):
with core.DeviceScope(muji.OnGPU(gpu_id)):
for blob_name in blob_names:
blob = core.ScopedName(blob_name)
workspace.CreateBlob(blob)
net.DequeueBlobs(
data_loader._blobs_queue_name, blob_names)
logger.info("Protobuf:\n" + str(net.Proto()))
return net
Example #14
| Source File: model_loader.py From VMZ with Apache License 2.0 | 6 votes |
|
def BroacastParameters(model, src_gpu, gpus):
log.info("Broadcasting parameters from gpu {} to gpu: {}".format(
src_gpu, ','.join([str(g) for g in gpus]))
)
for param in model.params:
if 'gpu_{}'.format(gpus[0]) in str(param):
for i in gpus:
blob = workspace.FetchBlob(str(param))
target_blob_name = str(param).replace(
'gpu_{}'.format(src_gpu),
'gpu_{}'.format(i)
)
log.info('broadcast {} -> {}'.format(
str(param), target_blob_name)
)
workspace.FetchBlob(str(param))
with core.DeviceScope(
core.DeviceOption(caffe2_pb2.CUDA, i)):
workspace.FeedBlob(target_blob_name, blob)
Example #15
| Source File: test_loader.py From Detectron-DA-Faster-RCNN with Apache License 2.0 | 6 votes |
|
def get_net(data_loader, name):
logger = logging.getLogger(__name__)
blob_names = data_loader.get_output_names()
net = core.Net(name)
net.type = 'dag'
for gpu_id in range(cfg.NUM_GPUS):
with core.NameScope('gpu_{}'.format(gpu_id)):
with core.DeviceScope(muji.OnGPU(gpu_id)):
for blob_name in blob_names:
blob = core.ScopedName(blob_name)
workspace.CreateBlob(blob)
net.DequeueBlobs(
data_loader._blobs_queue_name, blob_names)
logger.info("Protobuf:\n" + str(net.Proto()))
return net
Example #16
| Source File: test_batch_permutation_op.py From Detectron-DA-Faster-RCNN with Apache License 2.0 | 6 votes |
|
def _run_op_test(self, X, I, check_grad=False):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
op = core.CreateOperator('BatchPermutation', ['X', 'I'], ['Y'])
workspace.FeedBlob('X', X)
workspace.FeedBlob('I', I)
workspace.RunOperatorOnce(op)
Y = workspace.FetchBlob('Y')
if check_grad:
gc = gradient_checker.GradientChecker(
stepsize=0.1,
threshold=0.001,
device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
)
res, grad, grad_estimated = gc.CheckSimple(op, [X, I], 0, [0])
self.assertTrue(res, 'Grad check failed')
Y_ref = X[I]
np.testing.assert_allclose(Y, Y_ref, rtol=1e-5, atol=1e-08)
Example #17
| Source File: test_spatial_narrow_as_op.py From Detectron-DA-Faster-RCNN with Apache License 2.0 | 6 votes |
|
def _run_test(self, A, B, check_grad=False):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
op = core.CreateOperator('SpatialNarrowAs', ['A', 'B'], ['C'])
workspace.FeedBlob('A', A)
workspace.FeedBlob('B', B)
workspace.RunOperatorOnce(op)
C = workspace.FetchBlob('C')
if check_grad:
gc = gradient_checker.GradientChecker(
stepsize=0.005,
threshold=0.005,
device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
)
res, grad, grad_estimated = gc.CheckSimple(op, [A, B], 0, [0])
self.assertTrue(res, 'Grad check failed')
dims = C.shape
C_ref = A[:dims[0], :dims[1], :dims[2], :dims[3]]
np.testing.assert_allclose(C, C_ref, rtol=1e-5, atol=1e-08)
Example #18
| Source File: test_loader.py From Detectron with Apache License 2.0 | 6 votes |
|
def get_net(data_loader, name):
logger = logging.getLogger(__name__)
blob_names = data_loader.get_output_names()
net = core.Net(name)
net.type = 'dag'
for gpu_id in range(cfg.NUM_GPUS):
with core.NameScope('gpu_{}'.format(gpu_id)):
with core.DeviceScope(muji.OnGPU(gpu_id)):
for blob_name in blob_names:
blob = core.ScopedName(blob_name)
workspace.CreateBlob(blob)
net.DequeueBlobs(
data_loader._blobs_queue_name, blob_names)
logger.info("Protobuf:\n" + str(net.Proto()))
return net
Example #19
| Source File: test_batch_permutation_op.py From Detectron with Apache License 2.0 | 6 votes |
|
def _run_op_test(self, X, I, check_grad=False):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
op = core.CreateOperator('BatchPermutation', ['X', 'I'], ['Y'])
workspace.FeedBlob('X', X)
workspace.FeedBlob('I', I)
workspace.RunOperatorOnce(op)
Y = workspace.FetchBlob('Y')
if check_grad:
gc = gradient_checker.GradientChecker(
stepsize=0.1,
threshold=0.001,
device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
)
res, grad, grad_estimated = gc.CheckSimple(op, [X, I], 0, [0])
self.assertTrue(res, 'Grad check failed')
Y_ref = X[I]
np.testing.assert_allclose(Y, Y_ref, rtol=1e-5, atol=1e-08)
Example #20
| Source File: test_spatial_narrow_as_op.py From Detectron with Apache License 2.0 | 6 votes |
|
def _run_test(self, A, B, check_grad=False):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
op = core.CreateOperator('SpatialNarrowAs', ['A', 'B'], ['C'])
workspace.FeedBlob('A', A)
workspace.FeedBlob('B', B)
workspace.RunOperatorOnce(op)
C = workspace.FetchBlob('C')
if check_grad:
gc = gradient_checker.GradientChecker(
stepsize=0.005,
threshold=0.005,
device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
)
res, grad, grad_estimated = gc.CheckSimple(op, [A, B], 0, [0])
self.assertTrue(res, 'Grad check failed')
dims = C.shape
C_ref = A[:dims[0], :dims[1], :dims[2], :dims[3]]
np.testing.assert_allclose(C, C_ref, rtol=1e-5, atol=1e-08)
Example #21
| Source File: test_loader.py From Detectron-Cascade-RCNN with Apache License 2.0 | 6 votes |
|
def get_net(data_loader, name):
logger = logging.getLogger(__name__)
blob_names = data_loader.get_output_names()
net = core.Net(name)
net.type = 'dag'
for gpu_id in range(cfg.NUM_GPUS):
with core.NameScope('gpu_{}'.format(gpu_id)):
with core.DeviceScope(muji.OnGPU(gpu_id)):
for blob_name in blob_names:
blob = core.ScopedName(blob_name)
workspace.CreateBlob(blob)
net.DequeueBlobs(
data_loader._blobs_queue_name, blob_names)
logger.info("Protobuf:\n" + str(net.Proto()))
return net
Example #22
| Source File: test_batch_permutation_op.py From Detectron-Cascade-RCNN with Apache License 2.0 | 6 votes |
|
def _run_op_test(self, X, I, check_grad=False):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
op = core.CreateOperator('BatchPermutation', ['X', 'I'], ['Y'])
workspace.FeedBlob('X', X)
workspace.FeedBlob('I', I)
workspace.RunOperatorOnce(op)
Y = workspace.FetchBlob('Y')
if check_grad:
gc = gradient_checker.GradientChecker(
stepsize=0.1,
threshold=0.001,
device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
)
res, grad, grad_estimated = gc.CheckSimple(op, [X, I], 0, [0])
self.assertTrue(res, 'Grad check failed')
Y_ref = X[I]
np.testing.assert_allclose(Y, Y_ref, rtol=1e-5, atol=1e-08)
Example #23
| Source File: test_spatial_narrow_as_op.py From Detectron-Cascade-RCNN with Apache License 2.0 | 6 votes |
|
def _run_test(self, A, B, check_grad=False):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
op = core.CreateOperator('SpatialNarrowAs', ['A', 'B'], ['C'])
workspace.FeedBlob('A', A)
workspace.FeedBlob('B', B)
workspace.RunOperatorOnce(op)
C = workspace.FetchBlob('C')
if check_grad:
gc = gradient_checker.GradientChecker(
stepsize=0.005,
threshold=0.005,
device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
)
res, grad, grad_estimated = gc.CheckSimple(op, [A, B], 0, [0])
self.assertTrue(res, 'Grad check failed')
dims = C.shape
C_ref = A[:dims[0], :dims[1], :dims[2], :dims[3]]
np.testing.assert_allclose(C, C_ref, rtol=1e-5, atol=1e-08)
Example #24
| Source File: test_batch_permutation_op.py From masktextspotter.caffe2 with Apache License 2.0 | 6 votes |
|
def _run_op_test(self, X, I, check_grad=False):
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
op = core.CreateOperator('BatchPermutation', ['X', 'I'], ['Y'])
workspace.FeedBlob('X', X)
workspace.FeedBlob('I', I)
workspace.RunOperatorOnce(op)
Y = workspace.FetchBlob('Y')
if check_grad:
gc = gradient_checker.GradientChecker(
stepsize=0.1,
threshold=0.001,
device_option=core.DeviceOption(caffe2_pb2.CUDA, 0)
)
res, grad, grad_estimated = gc.CheckSimple(op, [X, I], 0, [0])
self.assertTrue(res, 'Grad check failed')
Y_ref = X[I]
np.testing.assert_allclose(Y, Y_ref, rtol=1e-5, atol=1e-08)
Example #25
| Source File: c2.py From Detectron-Cascade-RCNN with Apache License 2.0 | 5 votes |
|
def CpuScope():
"""Create a CPU device scope."""
cpu_dev = core.DeviceOption(caffe2_pb2.CPU)
with core.DeviceScope(cpu_dev):
yield
Example #26
| Source File: c2.py From masktextspotter.caffe2 with Apache License 2.0 | 5 votes |
|
def CpuScope():
"""Create a CPU device scope."""
cpu_dev = core.DeviceOption(caffe2_pb2.CPU)
with core.DeviceScope(cpu_dev):
yield
Example #27
| Source File: c2.py From Detectron-Cascade-RCNN with Apache License 2.0 | 5 votes |
|
def CudaScope(gpu_id):
"""Create a CUDA device scope for GPU device `gpu_id`."""
gpu_dev = CudaDevice(gpu_id)
with core.DeviceScope(gpu_dev):
yield
Example #28
| Source File: c2.py From masktextspotter.caffe2 with Apache License 2.0 | 5 votes |
|
def CudaScope(gpu_id):
"""Create a CUDA device scope for GPU device `gpu_id`."""
gpu_dev = CudaDevice(gpu_id)
with core.DeviceScope(gpu_dev):
yield
Example #29
| Source File: segmentation_no_db_example.py From peters-stuff with GNU General Public License v3.0 | 5 votes |
|
def add_training_operators(output_segmentation, model, device_opts) :
with core.DeviceScope(device_opts):
loss = model.SigmoidCrossEntropyWithLogits([output_segmentation, "gt_segmentation"], 'loss')
avg_loss = model.AveragedLoss(loss, "avg_loss")
model.AddGradientOperators([loss])
opt = optimizer.build_adam(model, base_learning_rate=0.01)
Example #30
| Source File: classification_no_db_example.py From peters-stuff with GNU General Public License v3.0 | 5 votes |
|
def create_model(m, device_opts) :
with core.DeviceScope(device_opts):
conv1 = brew.conv(m, 'data', 'conv1', dim_in=1, dim_out=20, kernel=5)
pool1 = brew.max_pool(m, conv1, 'pool1', kernel=2, stride=2)
conv2 = brew.conv(m, pool1, 'conv2', dim_in=20, dim_out=50, kernel=5)
pool2 = brew.max_pool(m, conv2, 'pool2', kernel=2, stride=2)
fc3 = brew.fc(m, pool2, 'fc3', dim_in=50 * 4 * 4, dim_out=500)
fc3 = brew.relu(m, fc3, fc3)
pred = brew.fc(m, fc3, 'pred', 500, 2)
softmax = brew.softmax(m, pred, 'softmax')
m.net.AddExternalOutput(softmax)
return softmax
# add loss and optimizer
