已合并
fix bug: fractal_z_3d格式转成ncdhw、ndhwc、dhwcn格式 #101
huafeng793创建于 12 天前
fix bug: fractal_z_3d格式转成ncdhw、ndhwc、dhwcn格式 #101
已合并
共 1 个文件变更+48-0
| @@ -356,6 +356,49 @@ def to_fractal_z_3d(data: numpy.ndarray, ori_format: str, target_shape: Union[li | |||
| 356 | return weight_group | 356 | return weight_group |
| 357 | 357 | ||
| 358 | 358 | ||
| 359 | +def from_fractal_z_3d(data: numpy.ndarray, target_shape: Union[list, tuple] = None, | ||
| 360 | + target_format: str = "NCDHW", groups=None): | ||
| 361 | + if groups is None: | ||
| 362 | + groups = 1 | ||
| 363 | + n = target_shape[target_format.index("N")] | ||
| 364 | + c_in = target_shape[target_format.index("C")] | ||
| 365 | + d = target_shape[target_format.index("D")] | ||
| 366 | + h = target_shape[target_format.index("H")] | ||
| 367 | + w = target_shape[target_format.index("W")] | ||
| 368 | + c0 = data.shape[-1] | ||
| 369 | + n0 = BLOCK_SIZE | ||
| 370 | + cin_ori = c_in | ||
| 371 | + cout_ori = n // groups | ||
| 372 | + | ||
| 373 | + group_dict = _calculate_group(c_in * groups, n, groups, c0) | ||
| 374 | + real_g = group_dict["real_g"] | ||
| 375 | + cin1_g = group_dict["cin1_g"] | ||
| 376 | + mag_factor = group_dict["mag_factor"] | ||
| 377 | + cout1_g = group_dict["cout1_g"] | ||
| 378 | + | ||
| 379 | + data = numpy.ascontiguousarray(data) | ||
| 380 | + data_reshaped = data.reshape((real_g, d, cin1_g, h, w, cout1_g, n0, c0)) | ||
| 381 | + | ||
| 382 | + result = numpy.zeros((n, c_in, d, h, w), dtype=data.dtype) | ||
| 383 | + for g in range(groups): | ||
| 384 | + for ci in range(cin_ori): | ||
| 385 | + for co in range(cout_ori): | ||
| 386 | + e = g % mag_factor | ||
| 387 | + dst_cin = e * cin_ori + ci | ||
| 388 | + dst_cout = e * cout_ori + co | ||
| 389 | + src_cout = g * cout_ori + co | ||
| 390 | + result[src_cout, ci, :, :, :] = data_reshaped[ | ||
| 391 | + g // mag_factor, :, dst_cin // c0, :, :, | ||
| 392 | + dst_cout // n0, dst_cout % n0, dst_cin % c0 | ||
| 393 | + ] | ||
| 394 | + | ||
| 395 | + if target_format == "NDHWC": | ||
| 396 | + return result.transpose(0, 2, 3, 4, 1) | ||
| 397 | + elif target_format == "DHWCN": | ||
| 398 | + return result.transpose(2, 3, 4, 1, 0) | ||
| 399 | + return result | ||
| 400 | + | ||
| 401 | + | ||
| 359 | def to_NC1HWC0(data: numpy.ndarray, ori_format: str, | 402 | def to_NC1HWC0(data: numpy.ndarray, ori_format: str, |
| 360 | target_shape: Union[list, tuple] = None): | 403 | target_shape: Union[list, tuple] = None): |
| 361 | ori_shape = data.shape | 404 | ori_shape = data.shape |
| @@ -496,5 +539,10 @@ format_transformation_map = { | |||
| 496 | "FRACTAL_NZ": { | 539 | "FRACTAL_NZ": { |
| 497 | "ND": nz2nd, | 540 | "ND": nz2nd, |
| 498 | }, | 541 | }, |
| 542 | + "FRACTAL_Z_3D": { | ||
| 543 | + "NCDHW": from_fractal_z_3d, | ||
| 544 | + "NDHWC": from_fractal_z_3d, | ||
| 545 | + "DHWCN": from_fractal_z_3d, | ||
| 546 | + }, | ||
| 499 | } | 547 | } |
| 500 | 548 | ||