jonstreeterandClaude
36860b09dd
This fixes the ModuleNotFoundError issues that occurred during clean installations. Changes: - Embedded colormnet_source dependencies directly into repository as: - colormnet_model/ (model architecture) - colormnet_inference/ (inference pipeline) - colormnet_util/ (utility functions) - colormnet_dataset/ (data transforms and utilities) - Updated all internal imports to use new embedded paths - Removed external colormnet_source path dependency from __init__.py - Added graceful fallback for spatial_correlation_sampler CUDA extension - Fixed shape calculation bugs in PyTorch fallback implementation - Added missing dependencies (hickle, scikit-image) to auto-install The repository is now self-contained and can be installed without requiring the separate colormnet_source directory. Tested and verified: - Node imports work correctly - Colorization pipeline runs successfully - Model loads and processes video frames Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
435 lines
18 KiB
Python
435 lines
18 KiB
Python
"""
|
|
trainer.py - warpper and utility functions for network training
|
|
Compute loss, back-prop, update parameters, logging, etc.
|
|
"""
|
|
|
|
|
|
import os
|
|
import time
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
|
|
from colormnet_model.network import ColorMNet
|
|
|
|
from colormnet_model.losses import LossComputer
|
|
from colormnet_util.log_integrator import Integrator
|
|
from colormnet_util.image_saver import pool_pairs_221128_TransColorization
|
|
|
|
# val
|
|
from torch.utils.data import DataLoader
|
|
from colormnet_inference.data.mask_mapper import MaskMapper
|
|
from colormnet_inference.inference_core import InferenceCore
|
|
import torch.nn.functional as F
|
|
from PIL import Image
|
|
from colormnet_util.transforms import lab2rgb_transform_PIL, calculate_psnr
|
|
|
|
|
|
class ColorMNetTrainer:
|
|
def __init__(self, config, logger=None, save_path=None, local_rank=0, world_size=1, wandb = None):
|
|
self.wandb= wandb
|
|
self.config = config
|
|
self.num_frames = config['num_frames']
|
|
self.num_ref_frames = config['num_ref_frames']
|
|
self.deep_update_prob = config['deep_update_prob']
|
|
self.local_rank = local_rank
|
|
|
|
self.model = nn.parallel.DistributedDataParallel(
|
|
ColorMNet(config).cuda(),
|
|
device_ids=[local_rank], output_device=local_rank, broadcast_buffers=False, find_unused_parameters=True)
|
|
|
|
# Set up logger when local_rank=0
|
|
self.logger = logger
|
|
self.save_path = save_path
|
|
if logger is not None:
|
|
self.last_time = time.time()
|
|
self.logger.log_string('model_size', str(sum([param.nelement() for param in self.model.parameters()])))
|
|
self.train_integrator = Integrator(self.logger, distributed=True, local_rank=local_rank, world_size=world_size)
|
|
self.loss_computer = LossComputer(config)
|
|
|
|
self.train()
|
|
self.optimizer = optim.AdamW(filter(
|
|
lambda p: p.requires_grad, self.model.parameters()), lr=config['lr'], weight_decay=config['weight_decay'])
|
|
self.scheduler = optim.lr_scheduler.MultiStepLR(self.optimizer, config['steps'], config['gamma'])
|
|
if config['amp']:
|
|
self.scaler = torch.cuda.amp.GradScaler()
|
|
|
|
# Logging info
|
|
self.log_text_interval = config['log_text_interval']
|
|
self.log_image_interval = config['log_image_interval']
|
|
self.save_network_interval = config['save_network_interval']
|
|
self.save_checkpoint_interval = config['save_checkpoint_interval']
|
|
if config['debug']:
|
|
self.log_text_interval = self.log_image_interval = 1
|
|
|
|
self.best_psnr = 0
|
|
self.best_it = 0
|
|
|
|
def do_val(self, it=0, val_dataset=None):
|
|
self.model.module.eval()
|
|
self.val()
|
|
|
|
print('starting validation at it: %s'%(it))
|
|
val_loader = val_dataset.get_datasets()
|
|
|
|
config = {}
|
|
config['enable_long_term'] = True
|
|
config['max_mid_term_frames'] = 10
|
|
config['min_mid_term_frames'] = 5
|
|
config['max_long_term_elements'] = 10000
|
|
config['num_prototypes'] = 128
|
|
config['benchmark'] = False
|
|
config['flip'] = False
|
|
config['top_k'] = 30
|
|
config['mem_every'] = 5
|
|
config['deep_update_every'] = -1
|
|
config['hidden_dim'] = 64
|
|
|
|
avg_psnr = []
|
|
wb_frames = []
|
|
|
|
# with torch.no_grad():
|
|
for vid_reader in val_loader:
|
|
clip_psnr = []
|
|
loader = DataLoader(vid_reader, batch_size=1, shuffle=False, num_workers=2)
|
|
vid_name = vid_reader.vid_name
|
|
vid_length = len(loader)
|
|
# no need to count usage for LT if the video is not that long anyway
|
|
config['enable_long_term_count_usage'] = (
|
|
config['enable_long_term'] and
|
|
(vid_length
|
|
/ (config['max_mid_term_frames']-config['min_mid_term_frames'])
|
|
* config['num_prototypes'])
|
|
>= config['max_long_term_elements']
|
|
)
|
|
|
|
mapper = MaskMapper()
|
|
processor = InferenceCore(self.model.module, config=config)
|
|
first_mask_loaded = False
|
|
|
|
t0 = 0
|
|
for ti, data in enumerate(loader):
|
|
t0 += 1
|
|
if (t0-1) % 10 != 0: # Do not validate every frame for speed up training
|
|
continue
|
|
|
|
with torch.cuda.amp.autocast(enabled=not config['benchmark']):
|
|
rgb = data['rgb'].cuda()[0]
|
|
msk = data.get('mask')
|
|
msk = msk[:,1:3,:,:] if msk is not None else None
|
|
|
|
info = data['info']
|
|
frame = info['frame'][0]
|
|
shape = info['shape']
|
|
need_resize = info['need_resize'][0]
|
|
|
|
"""
|
|
For timing see https://discuss.pytorch.org/t/how-to-measure-time-in-pytorch/26964
|
|
Seems to be very similar in testing as my previous timing method
|
|
with two cuda sync + time.time() in STCN though
|
|
"""
|
|
|
|
if not first_mask_loaded:
|
|
if msk is not None:
|
|
first_mask_loaded = True
|
|
else:
|
|
# no point to do anything without a mask
|
|
continue
|
|
|
|
if config['flip']:
|
|
rgb = torch.flip(rgb, dims=[-1])
|
|
msk = torch.flip(msk, dims=[-1]) if msk is not None else None
|
|
|
|
# Map possibly non-continuous labels to continuous ones
|
|
if msk is not None:
|
|
msk = torch.Tensor(msk[0]).cuda()
|
|
if need_resize:
|
|
msk = vid_reader.resize_mask(msk.unsqueeze(0))[0]
|
|
|
|
processor.set_all_labels(list(range(1,3)))
|
|
labels = range(1,3)
|
|
else:
|
|
labels = None
|
|
|
|
# Run the model on this frame
|
|
# print('******************* START %s *************'%ti)
|
|
prob = processor.step(rgb, msk, labels, end=(ti==vid_length-1))
|
|
# print('******************* END %s *************'%ti)
|
|
|
|
# Upsample to original size if needed
|
|
if need_resize:
|
|
prob = F.interpolate(prob.unsqueeze(1), shape, mode='bilinear', align_corners=False)[:,0]
|
|
rgb = F.interpolate(rgb.unsqueeze(1), shape, mode='bilinear', align_corners=False)[:,0]
|
|
|
|
if config['flip']:
|
|
prob = torch.flip(prob, dims=[-1])
|
|
|
|
# Save the mask
|
|
if info['save'][0]:
|
|
out_mask_final = lab2rgb_transform_PIL(torch.cat([rgb[:1,:,:], prob], dim=0))
|
|
out_mask_final = out_mask_final * 255
|
|
out_mask_final = out_mask_final.astype(np.uint8)
|
|
|
|
out_img = np.array(Image.fromarray(out_mask_final))
|
|
|
|
gt_folder = self.config['validation_root']
|
|
gt_path = os.path.join(gt_folder, info['vid_name'][0], info['frame'][0])
|
|
gt_img = np.array(Image.open(gt_path))
|
|
psnr = calculate_psnr(gt_img, out_img)
|
|
# avg_psnr.append(psnr)
|
|
clip_psnr.append(psnr)
|
|
|
|
save_shape = (384, 384) # resize to save wandb space
|
|
out_img = np.array(Image.fromarray(out_img).resize(save_shape, resample=0))
|
|
gt_img = np.array(Image.fromarray(gt_img).resize(save_shape, resample=0))
|
|
|
|
wb_frames.append(self.wandb.Image(out_img, caption="Pred_%s_it%s"%(t0, it)))
|
|
wb_frames.append(self.wandb.Image(gt_img, caption="GT%s_it%s"%(t0, it)))
|
|
self.wandb.log({"val/pairs": wb_frames},step=it)
|
|
|
|
|
|
print('current item: %s clip_psnr is: %s'%(info['vid_name'][0], np.mean(clip_psnr)))
|
|
avg_psnr += clip_psnr
|
|
|
|
self.wandb.log({'val/psnr': np.mean(avg_psnr)}, step=it)
|
|
self.logger.log_scalar('val/psnr', np.mean(avg_psnr), it)
|
|
print('finish validation at it: %s avg_psnr: %s best_psnr: %s best_it: %s'%(it, np.mean(avg_psnr), self.best_psnr, self.best_it))
|
|
|
|
return np.mean(avg_psnr)
|
|
|
|
def do_pass(self, data, it=0, val_dataset=None):
|
|
self.model.module.train()
|
|
self.train()
|
|
|
|
# No need to store the gradient outside training
|
|
torch.set_grad_enabled(self._is_train)
|
|
|
|
for k, v in data.items():
|
|
if type(v) != list and type(v) != dict and type(v) != int:
|
|
data[k] = v.cuda(non_blocking=True)
|
|
|
|
out = {}
|
|
frames = data['rgb']
|
|
first_frame_gt = data['first_frame_gt'].float()
|
|
b = frames.shape[0]
|
|
num_filled_objects = [o.item() for o in data['info']['num_objects']]
|
|
num_objects = first_frame_gt.shape[2]
|
|
selector = data['selector'].unsqueeze(2).unsqueeze(2)
|
|
|
|
wb_frames = []
|
|
with torch.cuda.amp.autocast(enabled=self.config['amp']):
|
|
# image features never change, compute once
|
|
key, shrinkage, selection, f16, f8, f4 = self.model('encode_key', frames)
|
|
|
|
filler_one = torch.zeros(1, dtype=torch.int64)
|
|
hidden = torch.zeros((b, num_objects, self.config['hidden_dim'], *key.shape[-2:]))
|
|
|
|
v16, hidden = self.model('encode_value', frames[:,0], f16[:,0], hidden, first_frame_gt[:,0])
|
|
|
|
values = v16.unsqueeze(3) # add the time dimension
|
|
|
|
for ti in range(1, self.num_frames):
|
|
if ti <= self.num_ref_frames:
|
|
ref_values = values
|
|
ref_keys = key[:,:,:ti]
|
|
ref_shrinkage = shrinkage[:,:,:ti] if shrinkage is not None else None
|
|
else:
|
|
# pick num_ref_frames random frames
|
|
# this is not very efficient but I think we would
|
|
# need broadcasting in gather which we don't have
|
|
indices = [
|
|
torch.cat([filler_one, torch.randperm(ti-1)[:self.num_ref_frames-1]+1])
|
|
for _ in range(b)]
|
|
ref_values = torch.stack([
|
|
values[bi, :, :, indices[bi]] for bi in range(b)
|
|
], 0)
|
|
ref_keys = torch.stack([
|
|
key[bi, :, indices[bi]] for bi in range(b)
|
|
], 0)
|
|
ref_shrinkage = torch.stack([
|
|
shrinkage[bi, :, indices[bi]] for bi in range(b)
|
|
], 0) if shrinkage is not None else None
|
|
|
|
# Segment frame ti
|
|
memory_readout = self.model('read_memory', key[:,:,ti], selection[:,:,ti] if selection is not None else None,
|
|
ref_keys, ref_shrinkage, ref_values)
|
|
|
|
# short term memory
|
|
memory_readout_short = self.model('read_memory_short', key[:,:,ti], key[:,:,ti-1], values[:, :, :, ti-1])
|
|
memory_readout += memory_readout_short
|
|
|
|
hidden, logits, masks = self.model('segment', (f16[:,ti], f8[:,ti], f4[:,ti]), memory_readout,
|
|
hidden, selector, h_out=(ti < (self.num_frames-1)))
|
|
|
|
# No need to encode the last frame
|
|
if ti < (self.num_frames-1):
|
|
is_deep_update = np.random.rand() < self.deep_update_prob
|
|
v16, hidden = self.model('encode_value', frames[:,ti], f16[:,ti], hidden, masks, is_deep_update=is_deep_update)
|
|
values = torch.cat([values, v16.unsqueeze(3)], 3)
|
|
|
|
out[f'masks_{ti}'] = masks
|
|
out[f'logits_{ti}'] = logits
|
|
|
|
if self._is_train:
|
|
if it % self.log_image_interval == 0 and it != 0:
|
|
if self.logger is not None:
|
|
|
|
out_mask_final = lab2rgb_transform_PIL(torch.cat([frames[0,ti,:1,:,:], masks[0]], dim=0))
|
|
out_mask_final = out_mask_final * 255
|
|
out_mask_final = out_mask_final.astype(np.uint8)
|
|
out_img = np.array(Image.fromarray(out_mask_final))
|
|
|
|
gt_image_final = lab2rgb_transform_PIL(torch.cat([frames[0,ti,:1,:,:], data['cls_gt'][0,ti]], dim=0))
|
|
gt_image_final = gt_image_final * 255
|
|
gt_image_final = gt_image_final.astype(np.uint8)
|
|
gt_img = np.array(Image.fromarray(gt_image_final))
|
|
|
|
wb_frames.append(self.wandb.Image(out_img, caption="Pred_%s"%ti))
|
|
wb_frames.append(self.wandb.Image(gt_img, caption="GT_%s"%ti))
|
|
|
|
|
|
if self._do_log or self._is_train:
|
|
losses = self.loss_computer.compute_l1loss({**data, **out}, num_filled_objects, it)
|
|
|
|
# Logging
|
|
if self._do_log:
|
|
self.integrator.add_dict(losses)
|
|
|
|
self.wandb.log({'train/loss': losses['total_loss'].item()}, step=it)
|
|
self.wandb.log({'train/dice_loss_7': losses['dice_loss_7'].item()}, step=it)
|
|
self.wandb.log({'train/lr': self.scheduler.get_last_lr()[0]}, step=it)
|
|
|
|
if self._is_train:
|
|
if it % self.log_image_interval == 0 and it != 0:
|
|
if self.logger is not None:
|
|
images = {**data, **out}
|
|
size = (384, 384) # resize to save wandb space
|
|
self.logger.log_cv2('train/pairs', pool_pairs_221128_TransColorization(images, size, num_filled_objects), it)
|
|
|
|
self.wandb.log({"train/pairs": wb_frames},step=it)
|
|
|
|
if self._is_train:
|
|
if (it) % self.log_text_interval == 0 and it != 0:
|
|
if self.logger is not None:
|
|
self.logger.log_scalar('train/lr', self.scheduler.get_last_lr()[0], it)
|
|
self.logger.log_metrics('train', 'time', (time.time()-self.last_time)/self.log_text_interval, it)
|
|
self.last_time = time.time()
|
|
self.train_integrator.finalize('train', it)
|
|
self.train_integrator.reset_except_hooks()
|
|
|
|
if it % self.save_network_interval == 0 and it != 0 and it >= 129999:
|
|
if self.logger is not None:
|
|
self.save_network(it)
|
|
|
|
if it % self.save_checkpoint_interval == 0 and it != 0:
|
|
if self.logger is not None:
|
|
self.save_checkpoint(it)
|
|
|
|
# Backward pass
|
|
self.optimizer.zero_grad(set_to_none=True)
|
|
if self.config['amp']:
|
|
self.scaler.scale(losses['total_loss']).backward()
|
|
self.scaler.step(self.optimizer)
|
|
self.scaler.update()
|
|
else:
|
|
losses['total_loss'].backward()
|
|
self.optimizer.step()
|
|
|
|
self.scheduler.step()
|
|
|
|
# validation
|
|
if it % self.save_network_interval == 0: # log_text_interval
|
|
current_psnr = self.do_val(it, val_dataset=val_dataset)
|
|
|
|
if current_psnr >= self.best_psnr:
|
|
self.best_psnr = current_psnr
|
|
self.best_it = it
|
|
self.save_best_network(it)
|
|
|
|
|
|
def save_network(self, it):
|
|
if self.save_path is None:
|
|
print('Saving has been disabled.')
|
|
return
|
|
|
|
os.makedirs(os.path.dirname(self.save_path), exist_ok=True)
|
|
model_path = f'{self.save_path}_{it}.pth'
|
|
torch.save(self.model.module.state_dict(), model_path)
|
|
print(f'Network saved to {model_path}.')
|
|
|
|
def save_best_network(self, it):
|
|
if self.save_path is None:
|
|
print('Saving has been disabled.')
|
|
return
|
|
|
|
os.makedirs(os.path.dirname(self.save_path), exist_ok=True)
|
|
model_path = f'{self.save_path}_best.pth'
|
|
torch.save(self.model.module.state_dict(), model_path)
|
|
print(f'Network saved to {model_path}.')
|
|
|
|
def save_checkpoint(self, it):
|
|
if self.save_path is None:
|
|
print('Saving has been disabled.')
|
|
return
|
|
|
|
os.makedirs(os.path.dirname(self.save_path), exist_ok=True)
|
|
checkpoint_path = f'{self.save_path}_checkpoint_{it}.pth'
|
|
checkpoint = {
|
|
'it': it,
|
|
'network': self.model.module.state_dict(),
|
|
'optimizer': self.optimizer.state_dict(),
|
|
'scheduler': self.scheduler.state_dict()}
|
|
torch.save(checkpoint, checkpoint_path)
|
|
print(f'Checkpoint saved to {checkpoint_path}.')
|
|
|
|
def load_checkpoint(self, path):
|
|
# This method loads everything and should be used to resume training
|
|
map_location = 'cuda:%d' % self.local_rank
|
|
checkpoint = torch.load(path, map_location={'cuda:0': map_location})
|
|
|
|
it = checkpoint['it']
|
|
network = checkpoint['network']
|
|
optimizer = checkpoint['optimizer']
|
|
scheduler = checkpoint['scheduler']
|
|
|
|
map_location = 'cuda:%d' % self.local_rank
|
|
self.model.module.load_state_dict(network)
|
|
self.optimizer.load_state_dict(optimizer)
|
|
self.scheduler.load_state_dict(scheduler)
|
|
|
|
print('Network weights, optimizer states, and scheduler states loaded.')
|
|
|
|
return it
|
|
|
|
def load_network_in_memory(self, src_dict):
|
|
self.model.module.load_weights(src_dict)
|
|
print('Network weight loaded from memory.')
|
|
|
|
def load_network(self, path):
|
|
# This method loads only the network weight and should be used to load a pretrained model
|
|
map_location = 'cuda:%d' % self.local_rank
|
|
src_dict = torch.load(path, map_location={'cuda:0': map_location})
|
|
|
|
self.load_network_in_memory(src_dict)
|
|
print(f'Network weight loaded from {path}')
|
|
|
|
def train(self):
|
|
self._is_train = True
|
|
self._do_log = True
|
|
self.integrator = self.train_integrator
|
|
self.model.eval()
|
|
return self
|
|
|
|
def val(self):
|
|
self._is_train = False
|
|
self._do_log = True
|
|
self.model.eval()
|
|
return self
|
|
|
|
def test(self):
|
|
self._is_train = False
|
|
self._do_log = False
|
|
self.model.eval()
|
|
return self
|