1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
|
import numpy as np
from caffe2.python import workspace, cnn, core
from caffe2.python import timeout_guard
from caffe2.proto import caffe2_pb2
def init_model(self):
# if cudnn needs to be turned off, several other places
# need to be modified:
# 1. operators need to be constructed with engine option, like below:
# conv_blob = model.Conv(...engine=engine)
# 2. when launch model, opts['model_param']['engine'] = "" instead of "CUDNN"
# 2. caffe2_disable_implicit_engine_preference in operator.cc set to true
train_model = cnn.CNNModelHelper(
order="NCHW",
name="resnet",
use_cudnn=False,
cudnn_exhaustive_search=False,
)
self.train_model = train_model
# test_model = cnn.CNNModelHelper(
# order="NCHW",
# name="resnet_test",
# use_cudnn=False,
# cudnn_exhaustive_search=False,
# init_params=False,
# )
self.test_model = None
self.log.info("Model creation completed")
def fun_per_epoch_b4RunNet(self, epoch):
pass
def fun_per_iter_b4RunNet(self, epoch, epoch_iter):
learning_rate = 0.05
for idx in range(self.opts['distributed']['first_xpu_id'],
self.opts['distributed']['first_xpu_id'] +
self.opts['distributed']['num_xpus']):
caffe2_pb2_device = caffe2_pb2.CUDA if \
self.opts['distributed']['device'] == 'gpu' else \
caffe2_pb2.CPU
with core.DeviceScope(core.DeviceOption(caffe2_pb2_device, idx)):
workspace.FeedBlob(
'{}_{}/lr'.format(self.opts['distributed']['device'], idx),
np.array(learning_rate, dtype=np.float32)
)
def run_training_net(self):
timeout = 2000.0
with timeout_guard.CompleteInTimeOrDie(timeout):
workspace.RunNet(self.train_model.net.Proto().name)
|
