def GetFunctionConfig(args): config = flow.function_config() config.enable_auto_mixed_precision(args.use_fp16) config.train.num_gradient_accumulation_steps(args.num_accumulation_steps) if args.use_xla: config.use_xla_jit(True) config.enable_fuse_add_to_output(True) config.enable_fuse_model_update_ops(True) return config
def _default_config(args): config = flow.function_config() config.default_logical_view(flow.scope.consistent_view()) config.default_data_type(flow.float) if args.use_fp16: config.enable_auto_mixed_precision(True) if args.use_xla: config.use_xla_jit(True) config.enable_fuse_add_to_output(True) return config
def _make_func_config(args): func_cfg = flow.function_config() if args.fp16: func_cfg.enable_auto_mixed_precision(True) func_cfg.prune_parallel_cast_ops(True) func_cfg.enable_fuse_add_to_output(True) func_cfg.enable_fuse_model_update_ops(True) func_cfg.enable_fuse_cast_scale(True) # turn on this flag when match ZeRO & DeepSpeed func_cfg.enable_non_distributed_optimizer(False) if args.num_accumulation_steps > 1: if hasattr(func_cfg.train, "num_gradient_accumulation_steps"): func_cfg.train.num_gradient_accumulation_steps( args.num_accumulation_steps) else: args.num_accumulation_steps = 1 print( "WARNING: This version of OneFlow dose not support gradient accumulation" " please try newer version.") return func_cfg
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. """ import unittest import numpy as np import oneflow.compatible.single_client.unittest from oneflow.compatible import single_client as flow config = flow.function_config() def make_job(input_shape, permute, dtype=flow.float32): config.use_xla_jit(False) config.use_tensorrt(False) @flow.global_function(config) def transpose_job(x=flow.FixedTensorDef(input_shape, dtype=dtype)): return flow.transpose(x, perm=permute) return transpose_job def make_xla_job(input_shape, permute, dtype=flow.float32): config.use_xla_jit(True)
} def load_image(image_path="test_img/ILSVRC2012_val_00020287.JPEG"): print(image_path) im = Image.open(image_path) im = im.resize((224, 224)) im = im.convert("RGB") # 有的图像是单通道的,不加转换会报错 im = np.array(im).astype("float32") im = (im - args.rgb_mean) / args.rgb_std im = np.transpose(im, (2, 0, 1)) im = np.expand_dims(im, axis=0) return np.ascontiguousarray(im, "float32") @flow.global_function("predict", flow.function_config()) def InferenceNet( images: tp.Numpy.Placeholder((1, 3, 224, 224), dtype=flow.float) ) -> tp.Numpy: logits = model_dict[args.model](images, args) predictions = flow.nn.softmax(logits) return predictions def main(): flow.env.log_dir(args.log_dir) assert os.path.isdir(args.model_load_dir) flow.load_variables(flow.checkpoint.get(args.model_load_dir)) image = load_image(args.image_path) predictions = InferenceNet(image) clsidx = predictions.argmax()