From 75df8a09cc1a6a0b5a9721d2a148979d256ffbde Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sun, 21 May 2023 13:52:42 +0000 Subject: [PATCH] upscaling pipeline --- custom_data/garment_uv_mask_full.npy | Bin 0 -> 131200 bytes data/dataset.py | 21 +- models/upscaler.py | 183 +++-------------- scripts/upscaler_train.py | 284 +++++++++++++++++++++++++++ utils/checkpoints.py | 7 +- 5 files changed, 322 insertions(+), 173 deletions(-) create mode 100644 custom_data/garment_uv_mask_full.npy create mode 100644 scripts/upscaler_train.py diff --git a/custom_data/garment_uv_mask_full.npy b/custom_data/garment_uv_mask_full.npy new file mode 100644 index 0000000000000000000000000000000000000000..c79c38e421db9565fadd773703d07501fa2523ac GIT binary patch literal 131200 zcmeI5y^bYG5runPo}#@5X`EeYF?awx5Do?w8rVS?h($992Y3oznEtzUE3+b_q9V?z z&O3d^>B@+VFU~o2yJvqiyZ`&`?|<{VKYY0R*WJG#e*FCP&tD#Xb^q|6KYjP`@&4g2 zKYjV{vRx=9}_S_cS_XvVeF;uo9+D@??V3B2^b$br^HpW zUCiD0MBnaoD^x+7&D3o#@Xb!YKvlF^P2Kj+{cpk!6}WGtm%!ZC?~0w7inc$DnA>80 z`&&#!r~2=WMbFd#IuXoW-Ky}Bsj2_Xn7J!#xVy;I^zQw40Wh=fUTgTs)YSg~WbO(Z z?k+Ml$$t|zGwNp1kpojy_x_oyj#zH)GDY2j|GV=7U~X#O>7fHtRr>*8ZaQSSxy)20 z|Es8(L9f;vePqgN-am8EQPa)+lZ2hLrm+?}FomuD4uH8P&yX`y4Sw|;Nc

(rhj3MHX!AK5BNBoUi+YR?gclF+t~Y?UOEh|M~+=ZHc{Xxm4&N)k!L zW}Vt|M4=?K?IT+yi6mmPPVG6OP!ihqk*$(M@~~N>qdA^Z4teAwTj!|cQL{crb3COS z^2kTF&QVFDW^In=Xi_QQv5#zrB9jKqx*X5Zq*B0RAK4B?CJUN1>CLgUGPs_PY>%On z#msv2=2%)8T+c_k#}LY4rY*XI9IXthYa`ucDCIEI7TrOPRtD9zk#91jGMIUX{v<;y zgXz1-cNtn4%)CQ?lA)Er^i5Q{1hE8U#i2h-5KBP%9x7deSOT)*(4QrUB_Mqdl`cUn z0a2@R6=Qit zbAfTJUNkLWpm+tX7|Sb~3yfp+qGPh^g|Zgopj5*+G1;(NFx@B+%vlVC zl8xfTRO4=;bfY*pXF(E5IEWKejk^ZYjl$rJ1xYC3AWl#98 zYarby49-}Pgc1(o1Xbg%fpnuVIAcK)N;rlS6b-zF(T$=2jRjFCNt4KnW-TC7=Y9fD%vwNNu~DuL~P3gMf#zqM>szR45} zpSz{o4`SD1^uldL5p{L=-t;jXc9a;HGlYbZ(#S+9lzfzi-#mnpdercnq<|6++vVT|J#M?uuk`ca^L9CSL66%m^eg>5 z_`F#T=5KJDg*$}}09)8>$KtoN&BC3+mH@46wnOO~*k)a%u%RFey#OS8%j&m<&RYzy z(klR>4OG7^bT&A+Qm+68#5JZR3E(@L<9-!5V0013CyDNHffPhx30swRn?XKv_0Rmc$3INYT zG`qqFk67t+do(9$~fU z*rNvp^sZt6up#2QCGj+*A{tjc2-=8o-I913K_T_49tds#xo$~34VaknRS$+|Q?6SQ zPc{-%UG#u>HfGt9d9ab7>LLilvoXt-%!7>tRTn`Zo{3qOBp*yDsCbbC!kL)Mj_g+x z3MyVSfp8|~vLpM|go26}O(2|!x$MY(HKCy5MH2{TVlF$fUri{e_^Js+voKlL>p6>y zsl9jqA6^2$glSUN?=kyqf-yz;e@Xz8rAf)Y^Yxqw#uQ)HU|1$)?IZteq7lVsNeGb5 z$UAN)Z91U(-u*xS1^~!$fN0-h{xrY?%1@dAm;oa1xSbr}0cH8Wh4;_W_)gJA0A_7xUSbvZVH_x_89yK!56y?PVh{*PJI zo!SH%`1PNybgwaewAIgCORVFi|L5Tjs45j1ETU^=c^_m?)I3E%Vk2MhBzwdbN}dOjOF& znptZVpQG7Xy;{0P9x8Qf&#bkI&(Z9xUJYF%3$?njXwp*6=V*FTvxcmZgIe8KG-;{k zb2L4vSwq&yL9uRZnzB^1IUJwTtRZS>px7{0O<5}19FI@wb|Y%6zuMWXnzB%JIi8== z?ncyDf3>q&HDjS_(wm>rZ?DPIU+wj~W;|Ctdh;{-?KOG&tG#~LjOVIHZ+=FO@m%%j&Clq!*W~H1_WE5jo~s`H`57HIntb`?ZnJB~L)oDZn9*{h$(LX5HoIm# zlpXqj87;3h`ReQ4^{yG8>kfUwjGot;eD(G2de@B4b%#D-M$c?TH8UmtE660L+CD(;0si z6cP`(B>!_J0C=?_#GLLwfd9zzI12fnGlB6N!*=F$ug^dBEstq>|8piVeq-3qobJ#2 z{^$KiOKQnt-?Bri$)rs?im@%4gBloX(xKI4lBONK_!h}S1&cT75Gz7SvyNJPi{znz z#e4LK1)-!-Pc6PjGEu2YjisT7MYG^>;| z=SrnCIX<#fibWTiRmz!jrBa$4A6Y8Jq6^I`<;=NKDNT-#ER|x>g=Up<=3J?iCdWsX zO0nodvr0L0u2f2s<0DI@SoEP;sgyZiDy2*Dk%dwm`p~RY%A7Bi(xv#wLMaY?XjUp^ z&X-E*Qha2g6o)=EE0rNt4KnW-TC7=Y9fD%vw zNNt4KnW-TC7=Y9fD%vw zNW zP6^>0>A)?K6z)R*FI}Lh(@|auFJks@oGu!@(6o46veUj^yy>UbWoxDBwJn{LfD%vw JN`R5T{{fnti6Q_1 literal 0 HcmV?d00001 diff --git a/data/dataset.py b/data/dataset.py index 188f1b2..e64dde9 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -30,14 +30,15 @@ def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file= if is_cape: frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0] + frames = [Frame(frame.split('.')[0], frame.split('.')[1]) for frame in frames] else: frames = [frame.split('.')[0] for frame in frames] # .split('_')[0] + frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames] - if testing: - frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') )) - frames = [frame.split('.')[0] + '_4' for frame in frames] + # if testing: + # frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') )) + # frames = [frame.split('.')[0] + '_4' for frame in frames] - frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames] if frames_file is not None: frames = [] @@ -65,12 +66,12 @@ class StaticCache(): if is_cape: self.uv_masks = dict() - uv_mask = np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32) - self.uv_masks[512] = uv_mask.astype(np.float32) + uv_mask = torch.tensor(np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32), device=device) + self.uv_masks[512] = uv_mask.to(torch.float32).cpu().numpy() for res in [32, 64, 128, 256]: - down = Resize(uv_mask[:,:,None], size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST) - down = np.squeeze(down.numpy().astype(np.float32)) - self.uv_masks[res] = down + down = Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST)(uv_mask[None,:,:]) + down = torch.squeeze(down).to(torch.float32) + self.uv_masks[res] = down.cpu().numpy() def read(self, id): @@ -127,7 +128,7 @@ class Dataset(Dataset): self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True) self.body_faces = np.load('smpl/smpl_faces.npy') - self.static_cache = StaticCache(folder, signature, is_cape=is_cape) + self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device) self.cache = dict() if use_cache else None diff --git a/models/upscaler.py b/models/upscaler.py index c03e663..b857ea3 100644 --- a/models/upscaler.py +++ b/models/upscaler.py @@ -1,172 +1,35 @@ -import tensorflow as tf -from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseWithSkip, ConvNeXtWithSkip +import torch +import torch.nn as nn - -class Bottleneck(tf.keras.Model): - - def __init__(self, dim, kernel_size=7): - super(Bottleneck, self).__init__() - - self.bottleneck = tf.keras.Sequential( - [Block(dim=dim, drop_path=0, kernel_size=kernel_size, - layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)] - ) - - def call(self, inputs, training): - x = inputs - - x = self.bottleneck(x, training=training) - - return x - - - -class Upscaler(tf.keras.Model): - - def __init__(self, bottleneck_size): +class Upscaler(nn.Module): + def __init__(self, low_res=32, kernel_size=3): super(Upscaler, self).__init__() - depths = [27, 27] - dims = [512, 512] - reverse_dims = [16] + dims[:-1] - self.encoder = ConvNeXtWithSkip(depths=depths, dims=dims, kernel_size=3) - - self.bottleneck = Bottleneck(dims[-1], kernel_size=3) - - self.decoder = ConvNeXtReverseWithSkip(depths=depths, dims=reverse_dims, kernel_size=3) - - def call(self, inputs, training): - img, mask = inputs - - img = tf.transpose(img, [0,2,3,1]) - - x, skip_features = self.encoder(img, training=training) - - x = self.bottleneck(x, training=training) - - x = self.decoder([x, skip_features], training=training) - - x = x * (1 - mask[:,:,:,tf.newaxis]) - x = x + img - x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5) - - return x - - -class UpscalerBasic(tf.keras.Model): - - def __init__(self, bottleneck_size, low_res=32, high_res=64): - super(UpscalerBasic, self).__init__() - - self.low_res = low_res - self.high_res = high_res - self.flatten = tf.keras.layers.Flatten() - - self.seq = tf.keras.Sequential([ - tf.keras.layers.Dense(1024, activation='relu'), - tf.keras.layers.BatchNormalization(), - tf.keras.layers.Dense(1024, activation='relu'), - tf.keras.layers.BatchNormalization(), - tf.keras.layers.Dense(1024, activation='relu'), - tf.keras.layers.BatchNormalization(), - tf.keras.layers.Dense(low_res * (low_res//2) * 3), - ]) - - self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3)) - - def call(self, inputs, training): - img, mask = inputs - - x = self.flatten(img) - x = self.seq(x, training=training) - x = self.reshape(x) - - x = x * (1 - mask[:,:,:,tf.newaxis]) - x = x + img - x = tf.image.resize(x, size=(self.high_res, self.high_res//2), method=tf.image.ResizeMethod.LANCZOS5) - - return x - - - - -class UpscalerBasicNew(tf.keras.Model): - - def __init__(self, low_res=32): - super(UpscalerBasicNew, self).__init__() - - self.low_res = low_res - self.flatten = tf.keras.layers.Flatten(name='upscale_enhancer_flattent') - - self.seq = tf.keras.Sequential([ - tf.keras.layers.Dense(4096, activation='relu', name='upscale_enhancer_dense_1'), - tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'), - tf.keras.layers.Dense(2048, activation='relu', name='upscale_enhancer_dense_2'), - tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'), - tf.keras.layers.Dense(2048, activation='relu', name='upscale_enhancer_dense_3'), - tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'), - tf.keras.layers.Dense(low_res * (low_res//2) * 3, activation='tanh', name='upscale_enhancer_dense_final'), - ]) - - self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3), name='upscale_enhancer_reshape') - - def call(self, inputs, training): - img, mask = inputs - - x = self.flatten(img) - x = self.seq(x, training=training) - x = self.reshape(x) - - x = x * (1 - mask[:,:,:,tf.newaxis]) - x = x + img - - return x - - -class UpscalerBasicConv(tf.keras.Model): - - def __init__(self, low_res=32, kernel_size=3): - super(UpscalerBasicConv, self).__init__() - self.low_res = low_res - self.seq = tf.keras.Sequential([ - tf.keras.layers.Conv2D(256, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_1'), - tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'), - tf.keras.layers.Conv2D(128, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_2'), - tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'), - tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_3'), - tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'), - tf.keras.layers.Conv2D(32, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_4'), - tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_4'), - tf.keras.layers.Dense(3, activation='tanh', name='upscale_enhancer_dense_final'), - ]) + self.seq = nn.Sequential( + nn.Conv2d(3, 256, kernel_size=kernel_size, padding='same'), + nn.ReLU(), + nn.BatchNorm2d(256), + nn.Conv2d(256, 128, kernel_size=kernel_size, padding='same'), + nn.ReLU(), + nn.BatchNorm2d(128), + nn.Conv2d(128, 64, kernel_size=kernel_size, padding='same'), + nn.ReLU(), + nn.BatchNorm2d(64), + nn.Conv2d(64, 32, kernel_size=kernel_size, padding='same'), + nn.ReLU(), + nn.BatchNorm2d(32), + nn.Conv2d(32, 3, kernel_size=kernel_size, padding='same'), + nn.Tanh() + ) - - def call(self, inputs, training): + def forward(self, inputs) -> torch.Tensor: img, mask = inputs - x = self.seq(img, training=training) + x = self.seq(img) - x = x * (1 - mask[:,:,:,tf.newaxis]) + x = x * (1 - mask.unsqueeze(1)) x = x + img return x - -# class Upscaler(tf.keras.Model): - -# def __init__(self, bottleneck_size): -# super(Upscaler, self).__init__() - -# self.decoder = ConvNeXtReverse(depths=list(reversed([9, 9, 9])), dims=list(reversed([bottleneck_size, 128, 32, 16]))) - -# def call(self, inputs, training): -# x, mask = inputs - -# x = tf.transpose(x, [0,2,3,1]) -# x = self.decoder(x) - -# x = x * mask[:,:,:,tf.newaxis] - -# return x - diff --git a/scripts/upscaler_train.py b/scripts/upscaler_train.py new file mode 100644 index 0000000..7abf6f5 --- /dev/null +++ b/scripts/upscaler_train.py @@ -0,0 +1,284 @@ +import os +import sys +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + +import importlib +import argparse +parser = argparse.ArgumentParser(description='') +parser.add_argument('--name', type=str) +parser.add_argument('--config', type=str) +parser.add_argument('--debug', action='store_true') +args = vars(parser.parse_args()) + +os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" +os.environ["CUDA_VISIBLE_DEVICES"] = '7' + + +import time +import argparse +import wandb +import random +import numpy as np +import torch +# from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM +from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, get_summary_str +import plotly.graph_objects as go +import plotly.express as ex +import torch.optim as optim +import torch.nn.functional as F + +from data.dataset import Dataset +# from torch.utils.data import DataLoader +from custom_utils.custom_utils import plot_3d +from models.upscaler import Upscaler +from utils.checkpoints import CheckpointSaver, load_model +from custom_utils.custom_utils import Accumulator +from utils.utils import InfiniteDataLoader as DataLoader + +from pytorch_msssim import MS_SSIM + + +low_res = 32 +is_skirt = False +is_cape = True +is_male = True +skirt_modifier = '_skirt' if is_skirt else '' +cape_modifier = '_cape' if is_cape else '' +gender_modifier = '' if not is_cape else ('_m' if is_male else '_f') +smpl_res = 3638 if is_cape else 14475 +use_cache = False +name=f'iccv_{low_res}' + skirt_modifier + cape_modifier + gender_modifier + +if not is_skirt: + train_dir = '/data-local/127/train' + valid_dir = '/data-local/127/valid' +else: + train_dir = '/data-local/128/train_skirt' + valid_dir = '/data-local/127/valid_skirt' + +if is_cape: + if is_male: + train_dir = '/data-local/127/cape_processed/train/male' + valid_dir = '/data-local/127/cape_processed/test/male' + else: + train_dir = '/data-local/127/cape_processed/train/female' + valid_dir = '/data-local/127/cape_processed/test/female' + + + +def init_wandb(): + run = wandb.init(project='upscale', name=name, notes='', dir='/data-local/127/temp', tags=['upscaler', 'iccv']) + artifact = wandb.Artifact('code', type='code') + artifact.add_file('scripts/upscaler_train.py') + artifact.add_file('losses/losses.py') + artifact.add_file('data/dataset.py') + config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')] + for config_file_name in config_file_names: + file_name = '/'.join(config_file_name.split('.')) + '.py' + if os.path.exists(file_name): + artifact.add_file(file_name) + run.log_artifact(artifact) + +uv2mesh = UV2Mesh(res=256, is_skirt=is_skirt, is_smpl_orig=is_cape) + + +def plot_3d(mesh_gt, mesh_pred_unposed, faces): + mesh_pred_unposed[:,2] += 1 + fig = go.Figure(data=[ + go.Mesh3d( + # 8 vertices of a cube + x=mesh_gt[:,0], + y=mesh_gt[:,2], + z=mesh_gt[:,1], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2] + ), + go.Mesh3d( + # 8 vertices of a cube + x=mesh_pred_unposed[:,0], + y=mesh_pred_unposed[:,2], + z=mesh_pred_unposed[:,1], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + ) + ]) + + fig.update_layout(scene_aspectmode='manual', + scene_aspectratio=dict(x=1, y=2, z=1)) + fig.update_layout( + scene = dict( + xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,3],), + zaxis = dict(nticks=4, range=[-1,1],),) + ) + + return fig + + +def train_step(batch): + model.train() + with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp): + pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]]) + pred = F.interpolate(pred, size=(512, 256), mode='bilinear') + + l1_loss = torch.mean(torch.sum(torch.abs(batch['uv_unposed'] - pred * batch['uv_mask'].unsqueeze(1)), dim=-1)) + if low_res == 32: + ms_ssim = torch.mean(_MSSSIM_small(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1)) + else: + ms_ssim = torch.mean(_MSSSIM(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1)) + + loss = 1e3 * (1 - ms_ssim) + + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad() + + model.eval() + with torch.inference_mode(): + mesh_pred = uv2mesh(pred) + mesh_orig = batch['mesh_unposed'] + metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'].unsqueeze(2) * (mesh_orig - mesh_pred) ** 2, dim=-1))) / torch.sum(batch['mesh_mask']) + + return { 'train/loss': loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, } + + +def valid_step(batch): + model.eval() + + with torch.inference_mode(): + pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]]) + pred = F.interpolate(pred, size=(512, 256), mode='bilinear') + + l1_loss = torch.mean(torch.sum(torch.abs(batch['uv_unposed'] - pred * batch['uv_mask'].unsqueeze(1)), dim=-1)) + if low_res == 32: + ms_ssim = torch.mean(_MSSSIM_small(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1)) + else: + ms_ssim = torch.mean(_MSSSIM(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1)) + + loss = 1e3 * (1 - ms_ssim) + mesh_pred = uv2mesh(pred) + mesh_orig = batch['mesh_unposed'] + metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'].unsqueeze(2) * (mesh_orig - mesh_pred) ** 2, dim=-1))) / torch.sum(batch['mesh_mask']) + + return { 'valid/loss': loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, } + + +def plot_everything(batch: dict, type: str): + model.eval() + with torch.inference_mode(): + pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]]) + pred = F.interpolate(pred, size=(512, 256), mode='bilinear') + + mesh_pred = uv2mesh(pred) + mesh_orig = batch['mesh_unposed'] + + _F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if is_cape else "") + ".npy") + mask = batch['mesh_mask'].cpu().numpy().astype(bool) + garment_unposed, faces = maskFacesBatch(mesh_pred.cpu().numpy(), _F, mask) + # gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])] + gt = [mesh_orig[i].cpu().numpy()[mask[i], :] for i in range(pred.shape[0])] + + figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i]) for i in range(4) } + # figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) } + + orig_img = batch['uv_unposed'] + img = torch.concat((orig_img, pred), axis=2) + imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)] + + return { **figs, **{ type + '/images': imgs} } + + +def train(): + global step + valid_iterator = iter(valid_dataloader) + + for batch in train_dataloader: + train_loss = train_step(batch) + train_accumulator.update(train_loss) + + if step % 25 == 0: + data_valid = next(valid_iterator) + valid_loss = valid_step(data_valid) + valid_accumulator.update(valid_loss) + + if step % 1000 == 0 or (step < 1000 and step % 1000 == 0): + to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } }) + if step % 5000 == 0 or (step < 5000 and step % 1000 == 0): + train_figs = plot_everything(batch, 'train') + valid_figs = plot_everything(data_valid, 'valid') + to_log = { **to_log, **train_figs, **valid_figs } + if step == 0: + to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } + + if not args['debug']: + wandb.log(to_log, step=step) + train_accumulator.reset() + valid_accumulator.reset() + print(step) + + if step % 10000 == 0: + ckpt_saver(model, step, optimizer, scaler) + + step += 1 + + +if __name__ == '__main__': + if not args['debug']: + init_wandb() + + use_amp = True + uv_mask_body = torch.ones((low_res, low_res//2), device='cuda') + low_resolution_modifier = '_' + str(low_res) + + signature = { 'id', 'frame_nr', + 'pose_shape', + 'uv_unposed' + low_resolution_modifier, + 'uv_unposed', + 'mesh_unposed', + 'mesh_posed', + 'mesh_mask', + 'uv_mask', + 'uv_mask' + low_resolution_modifier, + } + + # dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt') + dataset_train = Dataset(train_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=None, resolutions=[512, low_res], is_cape=is_cape) + dataset_valid = Dataset(valid_dir, device='cuda', use_cache=use_cache, signature=signature, frames_file=None, resolutions=[512, low_res], is_cape=is_cape) + train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True) + valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True) + train_accumulator = Accumulator() + valid_accumulator = Accumulator() + + if low_res == 32 or low_res == 64: + kernel_size = 3 + elif low_res == 128 or low_res == 256: + kernel_size = 5 + else: + kernel_size = 7 + + model : Upscaler = Upscaler(kernel_size=kernel_size).to('cuda') + optimizer = optim.Adam(model.parameters(), lr=0.0001) + scaler = torch.cuda.amp.GradScaler(enabled=use_amp) + + # _MSSSIM = MSSSIM(data_range=2).to(config['device']) + _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3, weights=[0.5, 0.2, 0.2, 0.1]) + _MSSSIM_small = MS_SSIM(data_range=2, size_average=True, channel=3, weights=[1.0]) + + step = 0 + ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + name, top_n=5) + + print(f'started training: {name}') + + if not args['debug']: + # model = torch.compile(model) + wandb.watch(model, log='all', log_freq=1000, log_graph=True) + + train() + + + diff --git a/utils/checkpoints.py b/utils/checkpoints.py index c098f44..07ebcae 100644 --- a/utils/checkpoints.py +++ b/utils/checkpoints.py @@ -16,9 +16,9 @@ class CheckpointSaver: self.top_n = top_n self.top_model_paths = [] - def __call__(self, model, iteration, optimizer): + def __call__(self, model, iteration, optimizer, scaler): model_path = os.path.join(self.dirpath, f'iteration_{iteration}.pt') - torch.save({ 'model': model.state_dict(), 'optimizer': optimizer }, model_path) + torch.save({ 'model': model.state_dict(), 'optimizer': optimizer, 'scaler': scaler }, model_path) self.top_model_paths.append({'path': model_path}) if len(self.top_model_paths)>self.top_n: self.cleanup() @@ -31,12 +31,13 @@ class CheckpointSaver: self.top_model_paths = self.top_model_paths[len(self.top_model_paths)-self.top_n:] -def load_model(model, checkpoint, optimizer=None, iteration=''): +def load_model(model, checkpoint, optimizer=None, scaler=None, iteration=''): # list of checkpoint files in the folder checkpoints checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'{iteration}.pt')] checkpoint = torch.load(checkpoints[-1]) iteration = checkpoints[-1].split('_')[-1].split('.')[0] model.load_state_dict(checkpoint['model']) if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer']) + if scaler is not None: scaler.load_state_dict(checkpoint['scaler']) return int(iteration) if iteration != '' else 0