prefacio
A veces es necesario guardar algunos mapas de características o resultados de salida de imágenes verificados durante el entrenamiento.
Sal a jugar más tarde, no escribas el prólogo
guardar_img
Esta manera es simple y grosera.
from torchvision.utils import save_image
save_image(tensor, path)
almohada.guardar
Primero convierta a numpy, vuelva a mapear de 0 a 255, luego convierta a almohada y finalmente guarde la imagen
def save_batch_img(batch_img, output_dir, prefix=None):
"""
img: tensor, [b, c, h, w] , [0, 1]
"""
for i in range(batch_img.shape[0]):
img = batch_img[i]
img = img.permute(1, 2, 0) # shape HWC
img = img.detach().cpu().numpy()
img = (img * 255).astype(np.uint8)
img = Image.fromarray(img)
# # 保存图像
os.makedirs(output_dir, exist_ok=True)
if prefix is not None:
img.save('{}/{}_{}.jpg'.format(output_dir, prefix, i))
else:
img.save('{}/{}.jpg'.format(output_dir, i))