Add segformer model. Start work on BYOL adaptation that will support training it.
This commit is contained in:
James Betker
parent
17555e7d07
commit
fc623d4b5a
327
codes/models/byol/byol_for_semantic_chaining.py
Normal file
327
codes/models/byol/byol_for_semantic_chaining.py
Normal file
|
|
@ -0,0 +1,327 @@
|
|||
import copy
|
||||
import os
|
||||
import random
|
||||
from functools import wraps
|
||||
import kornia.augmentation as augs
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torchvision
|
||||
from PIL import Image
|
||||
from kornia import filters, apply_hflip
|
||||
from torch import nn
|
||||
from torchvision.transforms import ToTensor
|
||||
|
||||
from data.byol_attachment import RandomApply
|
||||
from trainer.networks import register_model, create_model
|
||||
from utils.util import checkpoint, opt_get
|
||||
|
||||
|
||||
def default(val, def_val):
|
||||
return def_val if val is None else val
|
||||
|
||||
|
||||
def flatten(t):
|
||||
return t.reshape(t.shape[0], -1)
|
||||
|
||||
|
||||
def singleton(cache_key):
|
||||
def inner_fn(fn):
|
||||
@wraps(fn)
|
||||
def wrapper(self, *args, **kwargs):
|
||||
instance = getattr(self, cache_key)
|
||||
if instance is not None:
|
||||
return instance
|
||||
|
||||
instance = fn(self, *args, **kwargs)
|
||||
setattr(self, cache_key, instance)
|
||||
return instance
|
||||
|
||||
return wrapper
|
||||
|
||||
return inner_fn
|
||||
|
||||
|
||||
def get_module_device(module):
|
||||
return next(module.parameters()).device
|
||||
|
||||
|
||||
def set_requires_grad(model, val):
|
||||
for p in model.parameters():
|
||||
p.requires_grad = val
|
||||
|
||||
|
||||
# loss fn
|
||||
def loss_fn(x, y):
|
||||
x = F.normalize(x, dim=-1, p=2)
|
||||
y = F.normalize(y, dim=-1, p=2)
|
||||
return 2 - 2 * (x * y).sum(dim=-1)
|
||||
|
||||
|
||||
# exponential moving average
|
||||
class EMA():
|
||||
def __init__(self, beta):
|
||||
super().__init__()
|
||||
self.beta = beta
|
||||
|
||||
def update_average(self, old, new):
|
||||
if old is None:
|
||||
return new
|
||||
return old * self.beta + (1 - self.beta) * new
|
||||
|
||||
|
||||
def update_moving_average(ema_updater, ma_model, current_model):
|
||||
for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
|
||||
old_weight, up_weight = ma_params.data, current_params.data
|
||||
ma_params.data = ema_updater.update_average(old_weight, up_weight)
|
||||
|
||||
|
||||
# MLP class for projector and predictor
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, dim, projection_size, hidden_size=4096):
|
||||
super().__init__()
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(dim, hidden_size),
|
||||
nn.BatchNorm1d(hidden_size),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Linear(hidden_size, projection_size)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = flatten(x)
|
||||
return self.net(x)
|
||||
|
||||
|
||||
# A wrapper class for training against networks that do not collapse into a small-dimensioned latent.
|
||||
class StructuralMLP(nn.Module):
|
||||
def __init__(self, dim, projection_size, hidden_size=4096):
|
||||
super().__init__()
|
||||
b, c, h, w = dim
|
||||
flattened_dim = c * h // 4 * w // 4
|
||||
self.net = nn.Sequential(
|
||||
nn.Conv2d(c, c, kernel_size=3, padding=1, stride=2),
|
||||
nn.BatchNorm2d(c),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(c, c, kernel_size=3, padding=1, stride=2),
|
||||
nn.BatchNorm2d(c),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Flatten(),
|
||||
nn.Linear(flattened_dim, hidden_size),
|
||||
nn.BatchNorm1d(hidden_size),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Linear(hidden_size, projection_size)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.net(x)
|
||||
|
||||
|
||||
# a wrapper class for the base neural network
|
||||
# will manage the interception of the hidden layer output
|
||||
# and pipe it into the projecter and predictor nets
|
||||
class NetWrapper(nn.Module):
|
||||
def __init__(self, net, projection_size, projection_hidden_size, layer=-2, use_structural_mlp=False):
|
||||
super().__init__()
|
||||
self.net = net
|
||||
self.layer = layer
|
||||
|
||||
self.projector = None
|
||||
self.projection_size = projection_size
|
||||
self.projection_hidden_size = projection_hidden_size
|
||||
self.structural_mlp = use_structural_mlp
|
||||
|
||||
self.hidden = None
|
||||
self.hook_registered = False
|
||||
|
||||
def _find_layer(self):
|
||||
if type(self.layer) == str:
|
||||
modules = dict([*self.net.named_modules()])
|
||||
return modules.get(self.layer, None)
|
||||
elif type(self.layer) == int:
|
||||
children = [*self.net.children()]
|
||||
return children[self.layer]
|
||||
return None
|
||||
|
||||
def _hook(self, _, __, output):
|
||||
self.hidden = output
|
||||
|
||||
def _register_hook(self):
|
||||
layer = self._find_layer()
|
||||
assert layer is not None, f'hidden layer ({self.layer}) not found'
|
||||
handle = layer.register_forward_hook(self._hook)
|
||||
self.hook_registered = True
|
||||
|
||||
@singleton('projector')
|
||||
def _get_projector(self, hidden):
|
||||
if self.structural_mlp:
|
||||
projector = StructuralMLP(hidden.shape, self.projection_size, self.projection_hidden_size)
|
||||
else:
|
||||
_, dim = hidden.flatten(1,-1).shape
|
||||
projector = MLP(dim, self.projection_size, self.projection_hidden_size)
|
||||
return projector.to(hidden)
|
||||
|
||||
def get_representation(self, x):
|
||||
if self.layer == -1:
|
||||
return self.net(x)
|
||||
|
||||
if not self.hook_registered:
|
||||
self._register_hook()
|
||||
|
||||
unused = self.net(x)
|
||||
hidden = self.hidden
|
||||
self.hidden = None
|
||||
assert hidden is not None, f'hidden layer {self.layer} never emitted an output'
|
||||
return hidden
|
||||
|
||||
def forward(self, x):
|
||||
representation = self.get_representation(x)
|
||||
projector = self._get_projector(representation)
|
||||
projection = checkpoint(projector, representation)
|
||||
return projection
|
||||
|
||||
|
||||
class BYOL(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
net,
|
||||
image_size,
|
||||
hidden_layer=-2,
|
||||
projection_size=256,
|
||||
projection_hidden_size=4096,
|
||||
moving_average_decay=0.99,
|
||||
use_momentum=True,
|
||||
structural_mlp=False,
|
||||
do_augmentation=False # In DLAS this was intended to be done at the dataset level. For massive batch sizes
|
||||
# this can overwhelm the CPU though, and it becomes desirable to do the augmentations
|
||||
# on the GPU again.
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.online_encoder = NetWrapper(net, projection_size, projection_hidden_size, layer=hidden_layer,
|
||||
use_structural_mlp=structural_mlp)
|
||||
|
||||
self.do_aug = do_augmentation
|
||||
if self.do_aug:
|
||||
augmentations = [ \
|
||||
RandomApply(augs.ColorJitter(0.8, 0.8, 0.8, 0.2), p=0.8),
|
||||
augs.RandomGrayscale(p=0.2),
|
||||
augs.RandomHorizontalFlip(),
|
||||
RandomApply(filters.GaussianBlur2d((3, 3), (1.5, 1.5)), p=0.1),
|
||||
augs.RandomResizedCrop((self.cropped_img_size, self.cropped_img_size))]
|
||||
self.aug = nn.Sequential(*augmentations)
|
||||
self.use_momentum = use_momentum
|
||||
self.target_encoder = None
|
||||
self.target_ema_updater = EMA(moving_average_decay)
|
||||
|
||||
self.online_predictor = MLP(projection_size, projection_size, projection_hidden_size)
|
||||
|
||||
# get device of network and make wrapper same device
|
||||
device = get_module_device(net)
|
||||
self.to(device)
|
||||
|
||||
# send a mock image tensor to instantiate singleton parameters
|
||||
self.forward(torch.randn(2, 3, image_size, image_size, device=device),
|
||||
torch.randn(2, 3, image_size, image_size, device=device))
|
||||
|
||||
@singleton('target_encoder')
|
||||
def _get_target_encoder(self):
|
||||
target_encoder = copy.deepcopy(self.online_encoder)
|
||||
set_requires_grad(target_encoder, False)
|
||||
for p in target_encoder.parameters():
|
||||
p.DO_NOT_TRAIN = True
|
||||
return target_encoder
|
||||
|
||||
def reset_moving_average(self):
|
||||
del self.target_encoder
|
||||
self.target_encoder = None
|
||||
|
||||
def update_for_step(self, step, __):
|
||||
assert self.use_momentum, 'you do not need to update the moving average, since you have turned off momentum for the target encoder'
|
||||
assert self.target_encoder is not None, 'target encoder has not been created yet'
|
||||
update_moving_average(self.target_ema_updater, self.target_encoder, self.online_encoder)
|
||||
|
||||
def get_debug_values(self, step, __):
|
||||
# In the BYOL paper, this is made to increase over time. Not yet implemented, but still logging the value.
|
||||
return {'target_ema_beta': self.target_ema_updater.beta}
|
||||
|
||||
def visual_dbg(self, step, path):
|
||||
if self.do_aug:
|
||||
torchvision.utils.save_image(self.im1.cpu().float(), os.path.join(path, "%i_image1.png" % (step,)))
|
||||
torchvision.utils.save_image(self.im2.cpu().float(), os.path.join(path, "%i_image2.png" % (step,)))
|
||||
|
||||
def forward(self, image_one, image_two):
|
||||
if self.do_aug:
|
||||
image_one = self.aug(image_one)
|
||||
image_two = self.aug(image_two)
|
||||
# Keep copies on hand for visual_dbg.
|
||||
self.im1 = image_one.detach().copy()
|
||||
self.im2 = image_two.detach().copy()
|
||||
|
||||
online_proj_one = self.online_encoder(image_one)
|
||||
online_proj_two = self.online_encoder(image_two)
|
||||
|
||||
online_pred_one = self.online_predictor(online_proj_one)
|
||||
online_pred_two = self.online_predictor(online_proj_two)
|
||||
|
||||
with torch.no_grad():
|
||||
target_encoder = self._get_target_encoder() if self.use_momentum else self.online_encoder
|
||||
target_proj_one = target_encoder(image_one).detach()
|
||||
target_proj_two = target_encoder(image_two).detach()
|
||||
|
||||
loss_one = loss_fn(online_pred_one, target_proj_two.detach())
|
||||
loss_two = loss_fn(online_pred_two, target_proj_one.detach())
|
||||
|
||||
loss = loss_one + loss_two
|
||||
return loss.mean()
|
||||
|
||||
|
||||
class PointwiseAugmentor(nn.Module):
|
||||
def __init__(self, img_size=224):
|
||||
super().__init__()
|
||||
self.jitter = RandomApply(augs.ColorJitter(0.8, 0.8, 0.8, 0.2), p=0.8)
|
||||
self.gray = augs.RandomGrayscale(p=0.2)
|
||||
self.blur = RandomApply(filters.GaussianBlur2d((3, 3), (1.5, 1.5)), p=0.1)
|
||||
self.rrc = augs.RandomResizedCrop((img_size, img_size))
|
||||
|
||||
# Given a point in the *destination* image, returns the same point in the source image, given the kornia RRC params.
|
||||
def reverse_rrc(self, dest_point, params):
|
||||
dh, dw = params['dst'][:,2,1]-params['dst'][:,0,1], params['dst'][:,2,0] - params['dst'][:,0,0]
|
||||
sh, sw = params['src'][:,2,1]-params['src'][:,0,1], params['src'][:,2,0] - params['src'][:,0,0]
|
||||
scale_h, scale_w = sh.float() / dh.float(), sw.float() / dw.float()
|
||||
t, l = dest_point
|
||||
t = (t.float() * scale_h).int()
|
||||
l = (l.float() * scale_w).int()
|
||||
return t + params['src'][:,0,1], l + params['src'][:,0,0]
|
||||
|
||||
def reverse_horizontal_flip(self, pt, input):
|
||||
t, l = pt
|
||||
center = input.shape[-1] // 2
|
||||
return t, 2 * center - l
|
||||
|
||||
def forward(self, x, points):
|
||||
d = self.jitter(x)
|
||||
d = self.gray(d)
|
||||
will_flip = random.random() > .5
|
||||
if will_flip:
|
||||
d = apply_hflip(d)
|
||||
d = self.blur(d)
|
||||
params = self.rrc.generate_parameters(d.shape)
|
||||
d = self.rrc(d, params=params)
|
||||
|
||||
rev = self.reverse_rrc(points, params)
|
||||
if will_flip:
|
||||
rev = self.reverse_horizontal_flip(rev, x)
|
||||
|
||||
if __name__ == '__main__':
|
||||
p = PointwiseAugmentor(256)
|
||||
t = ToTensor()(Image.open('E:\\4k6k\\datasets\\ns_images\\imagesets\\000001_152761.jpg')).unsqueeze(0).repeat(8,1,1,1)
|
||||
points = (torch.randint(0,224,(t.shape[0],)),torch.randint(0,224,(t.shape[0],)))
|
||||
p(t, points)
|
||||
|
||||
|
||||
@register_model
|
||||
def register_byol(opt_net, opt):
|
||||
subnet = create_model(opt, opt_net['subnet'])
|
||||
return BYOL(subnet, opt_net['image_size'], opt_net['hidden_layer'],
|
||||
structural_mlp=opt_get(opt_net, ['use_structural_mlp'], False),
|
||||
do_augmentation=opt_get(opt_net, ['gpu_augmentation'], False))
|
||||
132
codes/models/segformer/backbone.py
Normal file
132
codes/models/segformer/backbone.py
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
# A direct copy of torchvision's resnet.py modified to support gradient checkpointing.
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torchvision.models.resnet import BasicBlock, Bottleneck
|
||||
from torchvision.models.utils import load_state_dict_from_url
|
||||
import torchvision
|
||||
|
||||
|
||||
__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
|
||||
'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',
|
||||
'wide_resnet50_2', 'wide_resnet101_2']
|
||||
|
||||
from trainer.networks import register_model
|
||||
from utils.util import checkpoint
|
||||
|
||||
model_urls = {
|
||||
'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
|
||||
'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
|
||||
'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
|
||||
'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
|
||||
'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
|
||||
'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
|
||||
'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
|
||||
'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
|
||||
'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
|
||||
}
|
||||
|
||||
|
||||
class Backbone(torchvision.models.resnet.ResNet):
|
||||
|
||||
def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
|
||||
groups=1, width_per_group=64, replace_stride_with_dilation=None,
|
||||
norm_layer=None):
|
||||
super().__init__(block, layers, num_classes, zero_init_residual, groups, width_per_group,
|
||||
replace_stride_with_dilation, norm_layer)
|
||||
del self.fc
|
||||
del self.avgpool
|
||||
|
||||
def _forward_impl(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
x = self.maxpool(x)
|
||||
|
||||
l1 = checkpoint(self.layer1, x)
|
||||
l2 = checkpoint(self.layer2, l1)
|
||||
l3 = checkpoint(self.layer3, l2)
|
||||
l4 = checkpoint(self.layer4, l3)
|
||||
|
||||
return l1, l2, l3, l4
|
||||
|
||||
def forward(self, x):
|
||||
return self._forward_impl(x)
|
||||
|
||||
|
||||
def _backbone(arch, block, layers, pretrained, progress, **kwargs):
|
||||
model = Backbone(block, layers, **kwargs)
|
||||
if pretrained:
|
||||
state_dict = load_state_dict_from_url(model_urls[arch],
|
||||
progress=progress)
|
||||
model.load_state_dict(state_dict)
|
||||
return model
|
||||
|
||||
|
||||
def backbone18(pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNet-18 model from
|
||||
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
return _backbone('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,
|
||||
**kwargs)
|
||||
|
||||
|
||||
def backbone34(pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNet-34 model from
|
||||
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
return _backbone('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,
|
||||
**kwargs)
|
||||
|
||||
|
||||
def backbone50(pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNet-50 model from
|
||||
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
return _backbone('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,
|
||||
**kwargs)
|
||||
|
||||
|
||||
def backbone101(pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNet-101 model from
|
||||
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
return _backbone('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress,
|
||||
**kwargs)
|
||||
|
||||
|
||||
def backbone152(pretrained=False, progress=True, **kwargs):
|
||||
r"""ResNet-152 model from
|
||||
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
|
||||
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
progress (bool): If True, displays a progress bar of the download to stderr
|
||||
"""
|
||||
return _backbone('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress,
|
||||
**kwargs)
|
||||
|
||||
|
||||
@register_model
|
||||
def register_resnet50(opt_net, opt):
|
||||
model = resnet50(pretrained=opt_net['pretrained'])
|
||||
if opt_net['custom_head_logits']:
|
||||
model.fc = nn.Linear(512 * 4, opt_net['custom_head_logits'])
|
||||
return model
|
||||
|
||||
89
codes/models/segformer/segformer.py
Normal file
89
codes/models/segformer/segformer.py
Normal file
|
|
@ -0,0 +1,89 @@
|
|||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from tqdm import tqdm
|
||||
|
||||
from models.segformer.backbone import backbone50
|
||||
|
||||
|
||||
class DilatorModule(nn.Module):
|
||||
def __init__(self, input_channels, output_channels, max_dilation):
|
||||
super().__init__()
|
||||
self.max_dilation = max_dilation
|
||||
self.conv1 = nn.Conv2d(input_channels, input_channels, kernel_size=3, padding=1, dilation=1, bias=True)
|
||||
if max_dilation > 1:
|
||||
self.bn = nn.BatchNorm2d(input_channels)
|
||||
self.relu = nn.ReLU()
|
||||
self.conv2 = nn.Conv2d(input_channels, input_channels, kernel_size=3, padding=1, dilation=max_dilation, bias=True)
|
||||
self.dense = nn.Linear(input_channels, output_channels, bias=True)
|
||||
|
||||
def forward(self, inp, loc):
|
||||
x = self.conv1(inp)
|
||||
if self.max_dilation > 1:
|
||||
x = self.bn(self.relu(x))
|
||||
x = self.conv2(x)
|
||||
|
||||
# This can be made (possibly substantially) more efficient by only computing these convolutions across a subset of the image. Possibly.
|
||||
i, j = loc
|
||||
x = x[:,:,i,j]
|
||||
return self.dense(x)
|
||||
|
||||
|
||||
# Grabbed from torch examples: https://github.com/pytorch/examples/tree/master/https://github.com/pytorch/examples/blob/master/word_language_model/model.py#L65:7
|
||||
class PositionalEncoding(nn.Module):
|
||||
def __init__(self, d_model, max_len=5000):
|
||||
super(PositionalEncoding, self).__init__()
|
||||
|
||||
pe = torch.zeros(max_len, d_model)
|
||||
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
|
||||
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
|
||||
pe[:, 0::2] = torch.sin(position * div_term)
|
||||
pe[:, 1::2] = torch.cos(position * div_term)
|
||||
pe = pe.unsqueeze(0).transpose(0, 1)
|
||||
self.register_buffer('pe', pe)
|
||||
|
||||
def forward(self, x):
|
||||
x = x + self.pe[:x.size(0), :]
|
||||
return x
|
||||
|
||||
|
||||
class Segformer(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.backbone = backbone50()
|
||||
backbone_channels = [256, 512, 1024, 2048]
|
||||
dilations = [[1,2,3,4],[1,2,3],[1,2],[1]]
|
||||
final_latent_channels = 2048
|
||||
dilators = []
|
||||
for ic, dis in zip(backbone_channels, dilations):
|
||||
layer_dilators = []
|
||||
for di in dis:
|
||||
layer_dilators.append(DilatorModule(ic, final_latent_channels, di))
|
||||
dilators.append(nn.ModuleList(layer_dilators))
|
||||
self.dilators = nn.ModuleList(dilators)
|
||||
|
||||
self.token_position_encoder = PositionalEncoding(final_latent_channels, max_len=10)
|
||||
self.transformer_layers = nn.Sequential(*[nn.TransformerEncoderLayer(final_latent_channels, nhead=4) for _ in range(16)])
|
||||
|
||||
def forward(self, x, pos):
|
||||
layers = self.backbone(x)
|
||||
set = []
|
||||
i, j = pos[0] // 4, pos[1] // 4
|
||||
for layer_out, dilator in zip(layers, self.dilators):
|
||||
for subdilator in dilator:
|
||||
set.append(subdilator(layer_out, (i, j)))
|
||||
i, j = i // 2, j // 2
|
||||
|
||||
# The torch transformer expects the set dimension to be 0.
|
||||
set = torch.stack(set, dim=0)
|
||||
set = self.token_position_encoder(set)
|
||||
set = self.transformer_layers(set)
|
||||
return set
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
model = Segformer().to('cuda')
|
||||
for j in tqdm(range(1000)):
|
||||
test_tensor = torch.randn(64,3,224,224).cuda()
|
||||
model(test_tensor, (43, 73))
|
||||
Loading…
Reference in New Issue
Block a user