Training crashed at epoch 1
Traceback for the error!
Traceback (most recent call last):
File “/data2/local_userdata/heyuting/CEU2/CEUTrack/lib/train/…/…/lib/train/trainers/base_trainer.py”, line 85, in train
self.train_epoch()
File “/data2/local_userdata/heyuting/CEU2/CEUTrack/lib/train/…/…/lib/train/trainers/ltr_trainer.py”, line 133, in train_epoch
self.cycle_dataset(loader)
File “/data2/local_userdata/heyuting/CEU2/CEUTrack/lib/train/…/…/lib/train/trainers/ltr_trainer.py”, line 86, in cycle_dataset
loss, stats = self.actor(data)
File “/data2/local_userdata/heyuting/CEU2/CEUTrack/lib/train/…/…/lib/train/actors/ceutrack.py”, line 285, in call
out_dict = self.forward_pass(data)
File “/data2/local_userdata/heyuting/CEU2/CEUTrack/lib/train/…/…/lib/train/actors/ceutrack.py”, line 334, in forward_pass
return_last_attn=False)
File “/home/heyuting/anaconda3/envs/ceu/lib/python3.7/site-packages/torch/nn/modules/module.py”, line 1130, in _call_impl
return forward_call(input, **kwargs)
File “/home/heyuting/anaconda3/envs/ceu/lib/python3.7/site-packages/torch/nn/parallel/distributed.py”, line 1008, in forward
output = self._run_ddp_forward(inputs, **kwargs)
File “/home/heyuting/anaconda3/envs/ceu/lib/python3.7/site-packages/torch/nn/parallel/distributed.py”, line 969, in _run_ddp_forward
return module_to_run(inputs[0], **kwargs[0])
File “/home/heyuting/anaconda3/envs/ceu/lib/python3.7/site-packages/torch/nn/modules/module.py”, line 1130, in _call_impl
return forward_call(input, **kwargs)
File “/data2/local_userdata/heyuting/CEU2/CEUTrack/lib/train/…/…/lib/models/ceutrack/ceutrack.py”, line 54, in forward
return_last_attn=return_last_attn)
File “/home/heyuting/anaconda3/envs/ceu/lib/python3.7/site-packages/torch/nn/modules/module.py”, line 1130, in _call_impl
return forward_call(input, **kwargs)
File “/data2/local_userdata/heyuting/CEU2/CEUTrack/lib/train/…/…/lib/models/ceutrack/vit_ce.py”, line 194, in forward
x, aux_dict = self.forward_features(z, x, event_z, event_x, ce_template_mask=ce_template_mask, ce_keep_rate=ce_keep_rate)
File “/data2/local_userdata/heyuting/CEU2/CEUTrack/lib/train/…/…/lib/models/ceutrack/vit_ce.py”, line 158, in forward_features
x = self.uep_module(x)
File “/home/heyuting/anaconda3/envs/ceu/lib/python3.7/site-packages/torch/nn/modules/module.py”, line 1130, in _call_impl
return forward_call(input, **kwargs)
File “/data2/local_userdata/heyuting/CEU2/CEUTrack/lib/train/…/…/lib/models/ceutrack/UEP.py”, line 37, in forward
assert N == self.num_patches and D == self.embed_dim, ‘Input shape does not match model configuration’
AssertionError: Input shape does not match model configuration

Restarting training from last epoch …
Finished training! 其中涉及到的文件有：① import math
import logging
from functools import partial
from collections import OrderedDict
from copy import deepcopy

import torch
import torch.nn as nn
import torch.nn.functional as F

from timm.models.layers import to_2tuple

from lib.models.layers.patch_embed import PatchEmbed, PatchEmbed_event, xcorr_depthwise
from .utils import combine_tokens, recover_tokens
from .vit import VisionTransformer
from …layers.attn_blocks import CEBlock
import random
import numpy as np
from .cross_module import Counter_Guide_Enhanced, Cross_Attention
from .UEP import UEPModule
_logger = logging.getLogger(name)


class VisionTransformerCE(VisionTransformer):
“”" Vision Transformer with candidate elimination (CE) module

A PyTorch impl of : An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
- https://arxiv.org/abs/2010.11929

Includes distillation token & head support for DeiT: Data-efficient Image Transformers
- https://arxiv.org/abs/2012.12877
“”"
def init(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
num_heads=12, mlp_ratio=4., qkv_bias=True, representation_size=None, distilled=False,
drop_rate=0., attn_drop_rate=0., drop_path_rate=0., embed_layer=PatchEmbed, norm_layer=None,
act_layer=None, weight_init=‘’,
ce_loc=None, ce_keep_ratio=None):
super().init()
if isinstance(img_size, tuple):
self.img_size = img_size
else:
self.img_size = to_2tuple(img_size)
self.patch_size = patch_size
self.in_chans = in_chans

self.num_classes = num_classes
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
self.num_tokens = 2 if distilled else 1
norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
act_layer = act_layer or nn.GELU

self.patch_embed = embed_layer(
img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
num_patches = self.patch_embed.num_patches

self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
self.pos_drop = nn.Dropout(p=drop_rate)
self.pos_embed_event = PatchEmbed_event(in_chans=32, embed_dim=768, kernel_size=4, stride=4)

dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
blocks = []
ce_index = 0
self.ce_loc = ce_loc
for i in range(depth):
ce_keep_ratio_i = 1.0
if ce_loc is not None and i in ce_loc:
ce_keep_ratio_i = ce_keep_ratio[ce_index]
ce_index += 1

blocks.append(
CEBlock(
dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, drop=drop_rate,
attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer,
keep_ratio_search=ce_keep_ratio_i)
)

self.blocks = nn.Sequential(blocks)
self.norm = norm_layer(embed_dim)

self.init_weights(weight_init)

# 添加交互模块counter_guide-sota模块
self.counter_guide = Counter_Guide_Enhanced(768, 768)
# 新增 UEP实例化
self.uep_module = UEPModule(num_patches,
embed_dim,
embed_dim // 4,
self.img_size)

def forward_features(self, z, x, event_z, event_x,
mask_z=None, mask_x=None,
ce_template_mask=None, ce_keep_rate=None,
return_last_attn=False
):
# 分支1 处理流程
B, H, W = x.shape[0], x.shape[2], x.shape[3]
x = self.patch_embed(x)
z = self.patch_embed(z)
z += self.pos_embed_z
x += self.pos_embed_x

event_x = self.pos_embed_event(event_x)
event_z = self.pos_embed_event(event_z)
event_x += self.pos_embed_x
event_z += self.pos_embed_z


if mask_z is not None and mask_x is not None:
mask_z = F.interpolate(mask_z[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0]
mask_z = mask_z.flatten(1).unsqueeze(-1)

mask_x = F.interpolate(mask_x[None].float(), scale_factor=1. / self.patch_size).to(torch.bool)[0]
mask_x = mask_x.flatten(1).unsqueeze(-1)

mask_x = combine_tokens(mask_z, mask_x, mode=self.cat_mode)
mask_x = mask_x.squeeze(-1)

if self.add_cls_token:
cls_tokens = self.cls_token.expand(B, -1, -1)
cls_tokens = cls_tokens + self.cls_pos_embed
if self.add_sep_seg:
x += self.search_segment_pos_embed
z += self.template_segment_pos_embed
x = combine_tokens(z, x, mode=self.cat_mode)
if self.add_cls_token:
x = torch.cat([cls_tokens, x], dim=1)
x = self.pos_drop(x)

lens_z = self.pos_embed_z.shape[1]
lens_x = self.pos_embed_x.shape[1]
global_index_t = torch.linspace(0, lens_z - 1, lens_z).to(x.device)
global_index_t = global_index_t.repeat(B, 1)
global_index_s = torch.linspace(0, lens_x - 1, lens_x).to(x.device)
global_index_s = global_index_s.repeat(B, 1)
removed_indexes_s = []

# # 分支2 处理流程
event_x = combine_tokens(event_z, event_x, mode=self.cat_mode)

if self.add_cls_token:
event_x = torch.cat([cls_tokens, event_x], dim=1)

lens_z = self.pos_embed_z.shape[1]
lens_x = self.pos_embed_x.shape[1]

global_index_t1 = torch.linspace(0, lens_z - 1, lens_z).to(event_x.device)
global_index_t1 = global_index_t1.repeat(B, 1)
global_index_s1 = torch.linspace(0, lens_x - 1, lens_x).to(event_x.device)
global_index_s1 = global_index_s1.repeat(B, 1)

removed_indexes_s1 = []

for i, blk in enumerate(self.blocks):
# 第一个分支处理
# 新增 uep 模块
x = self.uep_module(x)
x, global_index_t, global_index_s, removed_index_s, attn = <br/> blk(x, global_index_t, global_index_s, mask_x, ce_template_mask, ce_keep_rate)
# 第二个分支处理
# 新增 uep 模块
event_x = self.uep_module(event_x)
event_x, global_index_t1, global_index_s1, removed_index_s1, attn1 = <br/> blk(event_x, global_index_t1, global_index_s1, mask_x, ce_template_mask, ce_keep_rate)

if self.ce_loc is not None and i in self.ce_loc:
removed_indexes_s.append(removed_index_s)
removed_indexes_s1.append(removed_index_s1)

if i == 0 :
enhanced_x, enhanced_event_x = self.counter_guide(x, event_x)
x = x + enhanced_x
event_x = event_x + enhanced_event_x

# 应用LayerNorm归一化处理
x = self.norm(x)
event_x = self.norm(event_x)
x_cat = torch.cat([event_x, x], dim=1)
x = x_cat
aux_dict = {
“attn”: attn,
“attn1”: attn1,
“removed_indexes_s”: removed_indexes_s, # used for visualization
‘removed_indexes_s1’: removed_indexes_s1, # used for visualization
}
return x, aux_dict

def forward(self, z, x, event_z, event_x,
ce_template_mask=None, ce_keep_rate=None,
tnc_keep_rate=None,
return_last_attn=False):

x, aux_dict = self.forward_features(z, x, event_z, event_x, ce_template_mask=ce_template_mask, ce_keep_rate=ce_keep_rate)

return x, aux_dict


def _create_vision_transformer(pretrained=False, **kwargs):
model = VisionTransformerCE(**kwargs)

if pretrained:
if ‘npz’ in pretrained:
model.load_pretrained(pretrained, prefix=‘’)
else:
checkpoint = torch.load(pretrained, map_location=“cpu”)
missing_keys, unexpected_keys = model.load_state_dict(checkpoint[“model”], strict=False)
print('Load pretrained model from: ’ + pretrained)

return model


def vit_base_patch16_224_ce(pretrained=False, **kwargs):
“”" ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
“”“
model_kwargs = dict(
patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)
model = _create_vision_transformer(pretrained=pretrained, **model_kwargs)
return model

def vit_large_patch16_224_ce(pretrained=False, **kwargs):
“”” ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
“”“
model_kwargs = dict(
patch_size=16, embed_dim=1024, depth=24, num_heads=16, **kwargs)
model = _create_vision_transformer(pretrained=pretrained, **model_kwargs)
return model ② import numpy as np
import torch.nn as nn

from timm.models.layers import to_2tuple
import torch
import torch.nn.functional as F

class PatchEmbed(nn.Module):
“”” 2D Image to Patch Embedding
“”“

def init(self, img_size=256, patch_size=16, in_chans=3, embed_dim=768, norm_layer=False, flatten=True):
super().init()
img_size = to_2tuple(img_size)
patch_size = to_2tuple(patch_size)
self.img_size = img_size
self.patch_size = patch_size
self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) # 224 / 16 = 14
self.num_patches = self.grid_size[0] * self.grid_size[1]
self.flatten = flatten

self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
self.norm = nn.LayerNorm(embed_dim) if norm_layer else nn.Identity()

def forward(self, x):
# allow different input size
# B, C, H, W = x.shape
# _assert(H == self.img_size[0], f"Input image height ({H}) doesn’t match model ({self.img_size[0]}).”)
# _assert(W == self.img_size[1], f"Input image width ({W}) doesn’t match model ({self.img_size[1]}).")
x = self.proj(x)
if self.flatten:
x = x.flatten(2).transpose(1, 2) # BCHW -> BNC
x = self.norm(x)
return x


class PatchEmbed_event(nn.Module):
def init(self, in_chans=320, embed_dim=768, kernel_size=5, stride=1, flatten=True, norm_layer=False):
super().init()
self.pos_embedding = nn.Conv1d(in_channels=3, out_channels=16, kernel_size=1, stride=1)
self.in_chans = in_chans
self.flatten = flatten
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=kernel_size, stride=stride)
# self.proj2 = nn.Conv2d(embed_dim, embed_dim, kernel_size=1, stride=1)
self.norm = nn.LayerNorm(embed_dim) if norm_layer else nn.Identity()
# self.attn = SelfAttn(768, 4, 3072, 0.1, ‘relu’)

def forward(self, x):
# allow different input size
x = x.type(torch.cuda.FloatTensor)
xyz = self.pos_embedding(x.squeeze(dim=1)[:, :3, :])
xyz = F.relu(xyz)
x = torch.cat([xyz, x.squeeze(dim=1)[:, 3:, :]], dim=1)
B, N, C = x.shape # 1 1 19 10000
H = W = int(np.sqrt(NC//self.in_chans))
x = x.reshape(B, self.in_chans, H, W) # B 19 100 100
x = self.proj(x) # B 768 16 16
# x = self.proj2(x) # B 768 16 16

if self.flatten:
x = x.flatten(2).transpose(1, 2) # BCHW -> BNC 768256
# x = self.attn(src1=x, pos_src1=pos_embed)
x = self.norm(x)
return x

def xcorr_depthwise(x, kernel): # x, event_x
“”“depthwise cross correlation
“””
batch = kernel.size(0)
channel = kernel.size(2)
H = W = int(np.sqrt(x.size(1)))
x = x.reshape(1, batchchannel, H, W)
kernel = kernel.reshape(batchchannel, 1, H, W)
corr_weight = F.conv2d(x, kernel, groups=batchchannel)
# out = out.reshape(batch, HW, channel)
out = x * corr_weight
return out.reshape(batch, HW, channel)
③ import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import to_2tuple

class UEPModule(nn.Module):
def init(self, num_patches, embed_dim, hidden_dim, img_size):
super(UEPModule, self).init()
self.num_patches = num_patches
self.embed_dim = embed_dim
self.hidden_dim = hidden_dim
self.img_size = to_2tuple(img_size)
self.patch_size = int((self.img_size[0] * self.img_size[1]) ** 0.5 / num_patches ** 0.5)
# Recalculate H and W
self.H = self.img_size[0] // self.patch_size
self.W = self.img_size[1] // self.patch_size

# 1x1 Conv for Low-dimensional Projection
self.conv1 = nn.Conv2d(embed_dim, hidden_dim, kernel_size=1)

# Parallel 1x1 Conv and 1x1 DW-Conv
self.conv2 = nn.Conv2d(hidden_dim, hidden_dim, kernel_size=1)
self.dw_conv = nn.Conv2d(hidden_dim, hidden_dim, kernel_size=1, groups=hidden_dim) # Assuming DW-Conv is depthwise

# Additional 1x1 Conv before Activation
self.additional_conv = nn.Conv2d(hidden_dim, hidden_dim, kernel_size=1)

# 1x1 Conv for High-dimensional Projection
self.conv3 = nn.Conv2d(hidden_dim, embed_dim, kernel_size=1)

# Activation
self.gelu = nn.GELU()

def forward(self, x):
# Reshape (B,N,D) → (B,C,H,W)
B, N, D = x.shape
assert N == self.num_patches and D == self.embed_dim, ‘Input shape does not match model configuration’
x = x.permute(0, 2, 1).view(B, D, self.H, self.W)

# Low-dimensional projection
identity = x
x = self.conv1(x)

# Element-wise adding of parallel convolutions
conv2_out = self.conv2(x)
dw_conv_out = self.dw_conv(x)
x = conv2_out + dw_conv_out

# Additional 1x1 Conv and GELU
x = self.additional_conv(x)
x = self.gelu(x)

# High-dimensional projection
x = self.conv3(x)
x = x + identity # Add identity after high-dimensional projection

# Reshape back to (B,N,D)
x = x.view(B, self.embed_dim, -1).permute(0, 2, 1)
return x
，中文回答。