When I training the EfficientViT, I tried to change the size of the batch size to 64， and an error was raised.
RuntimeError: The size of tensor a (64) must match the size of tensor b (32) at non-singleton dimension 0 . In efficient-vit/efficient_vit.py line 169 ： x += self.pos_embedding[0:shape]. I think there may be an error in setting the location code.
I modified it. I don't know if it's correct, but it works normally
import torch
from torch import nn
from einops import rearrange
from efficientnet_pytorch import EfficientNet
import cv2
import re
from utils import resize
import numpy as np
from torch import einsum
from random import randint
class Residual(nn.Module):
def __init__(self, fn):
super().__init__()
self.fn = fn
def forward(self, x, **kwargs):
return self.fn(x, **kwargs) + x
class PreNorm(nn.Module):
def __init__(self, dim, fn):
super().__init__()
self.norm = nn.LayerNorm(dim)
self.fn = fn
def forward(self, x, **kwargs):
return self.fn(self.norm(x), **kwargs)
class FeedForward(nn.Module):
def __init__(self, dim, hidden_dim, dropout=0.):
super().__init__()
self.net = nn.Sequential(
nn.Linear(dim, hidden_dim),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, dim),
nn.Dropout(dropout)
)
def forward(self, x):
return self.net(x)
class Attention(nn.Module):
def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
super().__init__()
inner_dim = dim_head * heads
project_out = not (heads == 1 and dim_head == dim)
self.heads = heads
self.scale = dim_head ** -0.5
self.attend = nn.Softmax(dim=-1)
self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
self.to_out = nn.Sequential(
nn.Linear(inner_dim, dim),
nn.Dropout(dropout)
) if project_out else nn.Identity()
def forward(self, x):
b, n, _, h = *x.shape, self.heads
qkv = self.to_qkv(x).chunk(3, dim=-1)
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), qkv)
dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
attn = self.attend(dots)
out = einsum('b h i j, b h j d -> b h i d', attn, v)
out = rearrange(out, 'b h n d -> b n (h d)')
return self.to_out(out)
class Transformer(nn.Module):
def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
super().__init__()
self.layers = nn.ModuleList([])
for _ in range(depth):
self.layers.append(nn.ModuleList([
PreNorm(dim, Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout)),
PreNorm(dim, FeedForward(dim=dim, hidden_dim=mlp_dim, dropout=0))
]))
def forward(self, x):
for attn, ff in self.layers:
x = attn(x) + x
x = ff(x) + x
return x
class EfficientViT(nn.Module):
def __init__(self, config, channels=512, selected_efficient_net=0):
super().__init__()
image_size = config['model']['image-size']
patch_size = config['model']['patch-size']
num_classes = config['model']['num-classes']
dim = config['model']['dim']
depth = config['model']['depth']
heads = config['model']['heads']
mlp_dim = config['model']['mlp-dim']
emb_dim = config['model']['emb-dim']
dim_head = config['model']['dim-head']
dropout = config['model']['dropout']
emb_dropout = config['model']['emb-dropout']
assert image_size % patch_size == 0, 'image dimensions must be divisible by the patch size'
self.selected_efficient_net = selected_efficient_net
if selected_efficient_net == 0:
self.efficient_net = EfficientNet.from_pretrained('efficientnet-b0')
else:
self.efficient_net = EfficientNet.from_pretrained('efficientnet-b7')
checkpoint = torch.load("weights/final_999_DeepFakeClassifier_tf_efficientnet_b7_ns_0_23", map_location="cpu")
state_dict = checkpoint.get("state_dict", checkpoint)
self.efficient_net.load_state_dict({re.sub("^module.", "", k): v for k, v in state_dict.items()}, strict=False)
for i in range(0, len(self.efficient_net._blocks)):
for index, param in enumerate(self.efficient_net._blocks[i].parameters()):
if i >= len(self.efficient_net._blocks) - 3:
param.requires_grad = True
else:
param.requires_grad = False
self.num_patches = (7 // patch_size) ** 2
patch_dim = channels * patch_size ** 2
self.patch_size = patch_size
self.pos_embedding = nn.Parameter(torch.randn(1, self.num_patches + 1, dim))
self.patch_to_embedding = nn.Linear(patch_dim, dim)
self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
self.dropout = nn.Dropout(emb_dropout)
self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
self.to_cls_token = nn.Identity()
self.mlp_head = nn.Sequential(
nn.Linear(dim, mlp_dim),
nn.ReLU(),
nn.Linear(mlp_dim, num_classes)
)
def forward(self, img, mask=None):
p = self.patch_size
x = self.efficient_net.extract_features(img) # 1280x7x7
# x = self.features(img)
'''
for im in img:
image = im.cpu().detach().numpy()
image = np.transpose(image, (1,2,0))
cv2.imwrite("images/image"+str(randint(0,1000))+".png", image)
x_scaled = []
for idx, im in enumerate(x):
im = im.cpu().detach().numpy()
for patch_idx, patch in enumerate(im):
patch = (255*(patch - np.min(patch))/np.ptp(patch))
im[patch_idx] = patch
#cv2.imwrite("patches/patches_"+str(idx)+"_"+str(patch_idx)+".png", patch)
x_scaled.append(im)
x = torch.tensor(x_scaled).cuda()
'''
# x2 = self.features(img)
y = rearrange(x, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=p, p2=p)
# y2 = rearrange(x2, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = p, p2 = p)
y = self.patch_to_embedding(y)
cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
x = torch.cat((cls_tokens, y), 1)
shape = x.shape[0]
# x += self.pos_embedding[0:shape]
x += self.pos_embedding[:, :(self.num_patches + 1)]
x = self.dropout(x)
x = self.transformer(x)
x = self.to_cls_token(x[:, 0])
return self.mlp_head(x)
When I training the EfficientViT, I tried to change the size of the batch size to 64， and an error was raised. RuntimeError: The size of tensor a (64) must match the size of tensor b (32) at non-singleton dimension 0 . In efficient-vit/efficient_vit.py line 169 ： x += self.pos_embedding[0:shape]. I think there may be an error in setting the location code. I modified it. I don't know if it's correct, but it works normally