rectifier/stage0_model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

import timm
from timm.models.resnet import *
#from timm.models.convnext import *




########################################################################################
# unet decoder

class MyDecoderBlock(nn.Module):
	def __init__(
			self,
			in_channel,
			skip_channel,
			out_channel,
			scale=2,
	):
		super().__init__()
		self.scale=scale
		self.conv1 = nn.Sequential(
			nn.Conv2d(in_channel + skip_channel, out_channel, kernel_size=3, padding=1, bias=False),
			nn.BatchNorm2d(out_channel),
			nn.ReLU(inplace=True),
		)
		self.attention1 = nn.Identity()
		self.conv2 = nn.Sequential(
			nn.Conv2d(out_channel, out_channel, kernel_size=3, padding=1, bias=False),
			nn.BatchNorm2d(out_channel),
			nn.ReLU(inplace=True),
		)
		self.attention2 = nn.Identity()

	def forward(self, x, skip=None):
		x = F.interpolate(x, scale_factor=self.scale, mode='nearest')#, align_corners=False) #nearest bilinear
		if skip is not None:
			x = torch.cat([x, skip], dim=1)
			x = self.attention1(x)
		x = self.conv1(x)
		x = self.conv2(x)
		x = self.attention2(x)
		return x

class MyUnetDecoder(nn.Module):
	def __init__(
			self,
			in_channel,
			skip_channel,
			out_channel,
			scale=[2,2,2,2]
	):
		super().__init__()
		self.center = nn.Identity()

		i_channel = [in_channel, ] + out_channel[:-1]
		s_channel = skip_channel
		o_channel = out_channel
		block = [
			MyDecoderBlock(i, s, o, sc)
			for i, s, o, sc in zip(i_channel, s_channel, o_channel, scale)
		]
		self.block = nn.ModuleList(block)

	def forward(self, feature, skip):
		d = self.center(feature)
		decode = []
		for i, block in enumerate(self.block):
			if 0:
				print(i, d.shape, skip[i].shape if skip[i] is not None else 'none')
				print(block.conv1[0])
				print('')
			s = skip[i]
			d = block(d, s)
			decode.append(d)
		last = d
		return last, decode

#---------------------------------------------------
def encode_with_convnext(e, x):
	encode = []# x 256

	x = e.stem(x)  # 64
	x = e.stages[0](x); encode.append(x)  # 64
	x = e.stages[1](x); encode.append(x)  # 32
	x = e.stages[2](x); encode.append(x)  # 16
	x = e.stages[3](x); encode.append(x)  # 8
	return encode

def encode_with_resnet(e, x):
	encode = []# x 256

	x = e.conv1(x)
	x = e.bn1(x)
	x = e.act1(x)
	#x = e.maxpool(x)

	x = e.layer1(x); encode.append(x)  # 128
	x = e.layer2(x); encode.append(x)  # 64
	x = e.layer3(x); encode.append(x)  # 32
	x = e.layer4(x); encode.append(x)  # 16
	return encode

def F_cross_entropy(logits, truth, ignore_index=255):
	truth = truth.long()
	mask = truth != ignore_index
	if mask.sum() == 0:
		print('all pixels are ignored in bce !!!')
		return logits.sum() * 0

	#class_weights = torch.tensor([1.0, 10.0], device=logits.device)

	ce_loss= F.cross_entropy(
	logits, truth, ignore_index=ignore_index)
	return ce_loss


class Net(nn.Module):
	def __init__(self, pretrained=True, cfg=None):
		super(Net, self).__init__()
		self.output_type = ['infer', 'loss']
		self.register_buffer('D', torch.tensor(0))
		self.register_buffer('mean', torch.tensor([0.485, 0.456, 0.406]).reshape(1, 3, 1, 1))
		self.register_buffer('std', torch.tensor([0.229, 0.224, 0.225]).reshape(1, 3, 1, 1))

		#arch = 'convnext_tiny.fb_in22k'  # 'convnext_small.fb_in22k'
		# arch = 'resnet34.a3_in1k'
		arch = 'resnet18d.ra4_e3600_r224_in1k'

		encoder_dim = {
			'resnet18d.ra4_e3600_r224_in1k': [64, 128, 256, 512],
			'resnet34.a3_in1k': [64, 128, 256, 512],
			'convnext_tiny_in22k': [96, 192, 384, 768],
			'convnext_tiny.fb_in22k': [96, 192, 384, 768],
			'convnext_small.fb_in22k': [96, 192, 384, 768], #96, 192, 384, 768
			'convnext_base.fb_in22k': [128, 256, 512, 1024], #96, 192, 384, 768
			'resnet50d': [64, 256, 512, 1024, 2048, ],
		}[arch]
		decoder_dim = [256, 128, 64, 32]

		self.encoder = timm.create_model(
			model_name=arch, pretrained=pretrained, in_chans=3, num_classes=0, global_pool=''
		)
		self.decoder = MyUnetDecoder(
			in_channel=encoder_dim[-1],
			skip_channel=encoder_dim[:-1][::-1]+[0],
			out_channel=decoder_dim,
			scale = [2,2,2,2]
		)
		self.marker = nn.Conv2d(decoder_dim[-1], 13 + 1, kernel_size=1)  # lead name text
		self.orientation = nn.Linear(encoder_dim[-1], 8)

	def forward(self, batch):
		device = self.D.device
		image = batch['image'].to(device)

		B, _3_, H, W = image.shape
		x = image.float() / 255
		x = (x - self.mean) / self.std

		# ---------------------------------------
		e = self.encoder
		encode = encode_with_resnet(e, x)
		pooled = F.adaptive_avg_pool2d(encode[-1],1).reshape(B,-1)
		#[print(f'encode_{i}', e.shape) for i,e in enumerate(encode)]

		last, decode = self.decoder(
			feature=encode[-1], skip=encode[:-1][::-1]+[None]
		)
		#[print(f'decode_{i}', e.shape) for i,e in enumerate(decode)]
		#print('last', last.shape)

		marker = self.marker(last)
		orientation = self.orientation(pooled)

		output = {}
		if 'loss' in self.output_type:
			output['marker_loss'] = F_cross_entropy(
				marker, batch['marker'].to(device))
			output['orientation_loss'] = F_cross_entropy(
				orientation, batch['orientation'].to(device))

		# ----
		if 'infer' in self.output_type:
			output['marker'] = torch.softmax(marker,1)
			output['orientation'] = torch.softmax(orientation,1)

		return output



def run_check_net():
	H, W = 960, 1280
	batch_size = 4

	batch = {
		'image' : torch.from_numpy(np.random.randint(0, 256, (batch_size, 3, H, W))).byte(),
		'marker': torch.from_numpy(np.random.choice(3, (batch_size, H, W))).byte(),
		'orientation': torch.from_numpy(np.random.choice(5, (batch_size))).byte(),
	}

	net = Net(pretrained=True).cuda()
	# print(net)

	with torch.no_grad():
		with torch.amp.autocast('cuda'):
			output = net(batch)
	# ---

	print('batch')
	for k, v in batch.items():
		print(f'{k:>32} : {v.shape} ')

	print('output')
	for k, v in output.items():
		if 'loss' not in k:
			print(f'{k:>32} : {v.shape} ')

	print('loss')
	for k, v in output.items():
		if 'loss' in k:
			print(f'{k:>32} : {v.item()} ')

# main #################################################################
if __name__ == '__main__':
	run_check_net()