So below is my pytorch model Temporal Shift Model for video classification.
temporal shift
import torch
import torch.nn as nn
class TemporalShift(nn.Module):
def __init__(self,n_segments=8,n_div=8):
super().__init__()
self.n_segment=n_segments
self.n_div=n_div
def forward(self,x):
return self.shift(x,self.n_segment,self.n_div)
@staticmethod
def shift(x,n_seg,fold_div):
nt,c,h,w=x.shape
batch=nt//n_seg
x=x.view(batch,n_seg,c,h,w)
fold=c//fold_div
out=torch.zeros_like(x)
#Forward shift
out[:,:-1,:fold]=x[:,1:,:fold]
#Backward shift
out[:,1:,fold:2*fold]=x[:,:-1,fold:2*fold]
#remaining unchanged
out[:,:,2*fold:]=x[:,:,2*fold:]
return out.view(nt,c,h,w)
wrapper to add this to resnet blocks
from .temporal_shift import *
import torch.nn as nn
import torch
class TSMBlock(nn.Module):
def __init__(self,block,n_segment=8,fold_div=8):
super().__init__()
self.tsm=TemporalShift(n_segments=n_segment,n_div=fold_div)
self.block=block
def forward(self,x):
x=self.tsm(x)
x=self.block(x)
return x
class ConsensusModule(nn.Module):
def __init__(self,consensus_type='avg',dim=1):
super().__init__()
self.consensus_type=consensus_type
self.dim=dim
def forward(self,x):
if self.consensus_type=='avg':
return x.mean(dim=self.dim)
elif self.consensus_type=='max':
return x.max(dim=self.dim)
elif self.consensus_type=='identity':
return x
else:
raise ValueError(f"Unknown Consensus Type:{self.consensus_type}")
Temporal Shift Resnet model
import torch
import torch.nn as nn
import torchvision.models as models
from .temporal_shift_block import *
class TSM(nn.Module):
def __init__(self,num_classes,backbone="resnet18",
num_segments=16,fold_div=8,pretrained=True,dropout=0.5,consensus='avg'):
"""
Parameters
----------
num_classes : int
no of classes.
backbone : TYPE, torch model
DESCRIPTION. The default is "resnet18".
num_segments : INT, optional
no of segments. The default is 16.
fold_div : TYPE, optional
DESCRIPTION. The default is 8.
pretrained : TYPE, optional
DESCRIPTION. The default is True.
dropout : TYPE, optional
DESCRIPTION. The default is 0.5.
consensus : TYPE, optional
DESCRIPTION. The default is 'avg'.
Raises
------
ValueError
DESCRIPTION.
Returns
-------
None.
"""
super().__init__()
self.n_segments=num_segments
#Backbone
if backbone=="resnet18":
self.backbone=models.resnet18(
weights=models.ResNet18_Weights.DEFAULT if pretrained
else None)
elif backbone=='resnet34':
self.backbone=models.resnet34(
weights=models.ResNet34_Weights.DEFAULT if pretrained
else None)
elif backbone=="resnet50":
self.backbone=models.resnet50(
weights=models.ResNet50_Weights.DEFAULT if pretrained
else None)
else:
raise ValueError(f"Unsupported Backbone:{backbone}")
self.insert_tsm(self.backbone,n_seg=self.n_segments,fold_div=fold_div)
feature_dim=self.backbone.fc.in_features
self.backbone.fc=nn.Identity()
self.consensus=ConsensusModule(consensus)
#Classifier
self.dropout=nn.Dropout(dropout)
self.classifier=nn.Linear(feature_dim,num_classes)
def insert_tsm(self,model,n_seg,fold_div):
"""
convert resnet block into tsm block
"""
for layer_name in["layer1","layer2","layer3","layer4"]:
layer=getattr(model,layer_name)
wrapped=[]
for block in layer:
wrapped.append(
TSMBlock(block,n_segment=n_seg,fold_div=fold_div)
)
setattr(model,layer_name,nn.Sequential(*wrapped))
def forward(self,x):
"""
x input of dim(B,T,C,H,W)
"""
B,C,T,H,W=x.shape
x=x.permute(0,2,1,3,4)
x=x.reshape(B*T,C,H,W)
#Feature extraction
x=self.backbone(x)
#(B*T,F)->(B,T,F)
x=x.view(B,T,-1)
#Temporal aggregation
x=self.consensus(x)
x=self.dropout(x)
logits=self.classifier(x)
return logits
if __name__=="__main__":
n_classes=3
tsm=TSM(n_classes,backbone='resnet50',num_segments=16,
fold_div=8,pretrained=True,dropout=0.5,consensus='avg')
Now model trains well without quantization but when i try to do fx quantization on it as follows
from models.tsm_resnet import *
from torch.ao.quantization.quantize_fx import (
prepare_fx,
convert_fx,
)
from drowsiness_dataloader import *
from torch.ao.quantization import get_default_qconfig_mapping
import numpy as np
import pandas as pd
import torch
import os
ckp='model_best_yawn.pth.tar'
batch_size=8
num_classes=4
from drowsiness_dataloader import YawnTSMDataset
train_vidoes,train_labels=load_data('train_lst/train.lst')
train_sample_weights=get_sample_weights(train_labels)
train_dataset=YawnTSMDataset(list_file='train_lst/train.lst',
num_segments=8,
is_train=True,
enable_crop=True)
train_sampler=WeightedRandomSampler(weights=train_sample_weights, num_samples=len(train_labels),
replacement=True)
train_loader=DataLoader(train_dataset,batch_size=batch_size,sampler=train_sampler,
shuffle=False,num_workers=0,pin_memory=True)
is_gpu = torch.cuda.is_available()
device = torch.device("cuda" if is_gpu else "cpu")
model=TSM(num_classes=num_classes,
backbone='resnet18',
num_segments=8,
fold_div=8,
pretrained=True,
dropout=0.5,
consensus='avg')
checkpoint=torch.load(ckp,map_location='cpu')
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()
batch = next(iter(train_loader))
example_images = batch["video"].cpu()
print("PyTorch:", torch.__version__)
print("Example input:", example_images.shape)
print("Model loaded successfully")
example_images = example_images.cpu()
print("\n" + "=" * 70)
print("FX GRAPH MODE QUANTIZATION")
print("=" * 70)
qconfig_mapping = get_default_qconfig_mapping("x86")
print("QConfig mapping created.")
# Prepare
prepared_model = prepare_fx(
model,
qconfig_mapping,
(example_images,),
)
prepared_model.eval()
print("FX prepare successful.")
i end up getting following error log
PyTorch: 2.1.2+cu121
Example input: torch.Size([8, 3, 8, 224, 224])
Model loaded successfully
======================================================================
FX GRAPH MODE QUANTIZATION
======================================================================
QConfig mapping created.
C:\Users\slvrd\anaconda3\Lib\site-packages\torch\overrides.py:110: UserWarning: 'has_cuda' is deprecated, please use 'torch.backends.cuda.is_built()'
torch.has_cuda,
C:\Users\slvrd\anaconda3\Lib\site-packages\torch\overrides.py:111: UserWarning: 'has_cudnn' is deprecated, please use 'torch.backends.cudnn.is_available()'
torch.has_cudnn,
C:\Users\slvrd\anaconda3\Lib\site-packages\torch\overrides.py:117: UserWarning: 'has_mps' is deprecated, please use 'torch.backends.mps.is_built()'
torch.has_mps,
C:\Users\slvrd\anaconda3\Lib\site-packages\torch\overrides.py:118: UserWarning: 'has_mkldnn' is deprecated, please use 'torch.backends.mkldnn.is_available()'
torch.has_mkldnn,
Traceback (most recent call last):
File ~\anaconda3\Lib\site-packages\spyder_kernels\py3compat.py:356 in compat_exec
exec(code, globals, locals)
File c:\users\slvrd\untitled5.py:73
prepared_model = prepare_fx(
File ~\anaconda3\Lib\site-packages\torch\ao\quantization\quantize_fx.py:382 in prepare_fx
return _prepare_fx(
File ~\anaconda3\Lib\site-packages\torch\ao\quantization\quantize_fx.py:135 in _prepare_fx
graph_module = GraphModule(model, tracer.trace(model))
File ~\anaconda3\Lib\site-packages\torch\fx\_symbolic_trace.py:817 in trace
(self.create_arg(fn(*args)),),
File ~\models\tsm_resnet.py:104 in forward
x=self.backbone(x)
File ~\anaconda3\Lib\site-packages\torch\fx\_symbolic_trace.py:795 in module_call_wrapper
return self.call_module(mod, forward, args, kwargs)
File ~\anaconda3\Lib\site-packages\torch\fx\_symbolic_trace.py:479 in call_module
ret_val = forward(*args, **kwargs)
File ~\anaconda3\Lib\site-packages\torch\fx\_symbolic_trace.py:788 in forward
return _orig_module_call(mod, *args, **kwargs)
File ~\anaconda3\Lib\site-packages\torch\nn\modules\module.py:1518 in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File ~\anaconda3\Lib\site-packages\torch\nn\modules\module.py:1527 in _call_impl
return forward_call(*args, **kwargs)
File ~\anaconda3\Lib\site-packages\torchvision\models\resnet.py:285 in forward
return self._forward_impl(x)
File ~\anaconda3\Lib\site-packages\torchvision\models\resnet.py:273 in _forward_impl
x = self.layer1(x)
File ~\anaconda3\Lib\site-packages\torch\fx\_symbolic_trace.py:795 in module_call_wrapper
return self.call_module(mod, forward, args, kwargs)
File ~\anaconda3\Lib\site-packages\torch\fx\_symbolic_trace.py:479 in call_module
ret_val = forward(*args, **kwargs)
File ~\anaconda3\Lib\site-packages\torch\fx\_symbolic_trace.py:788 in forward
return _orig_module_call(mod, *args, **kwargs)
File ~\anaconda3\Lib\site-packages\torch\nn\modules\module.py:1518 in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File ~\anaconda3\Lib\site-packages\torch\nn\modules\module.py:1527 in _call_impl
return forward_call(*args, **kwargs)
File ~\anaconda3\Lib\site-packages\torch\nn\modules\container.py:215 in forward
input = module(input)
File ~\anaconda3\Lib\site-packages\torch\fx\_symbolic_trace.py:795 in module_call_wrapper
return self.call_module(mod, forward, args, kwargs)
File ~\anaconda3\Lib\site-packages\torch\fx\_symbolic_trace.py:479 in call_module
ret_val = forward(*args, **kwargs)
File ~\anaconda3\Lib\site-packages\torch\fx\_symbolic_trace.py:788 in forward
return _orig_module_call(mod, *args, **kwargs)
File ~\anaconda3\Lib\site-packages\torch\nn\modules\module.py:1518 in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File ~\anaconda3\Lib\site-packages\torch\nn\modules\module.py:1527 in _call_impl
return forward_call(*args, **kwargs)
File ~\models\temporal_shift_block.py:23 in forward
x=self.tsm(x)
File ~\anaconda3\Lib\site-packages\torch\fx\_symbolic_trace.py:795 in module_call_wrapper
return self.call_module(mod, forward, args, kwargs)
File ~\anaconda3\Lib\site-packages\torch\fx\_symbolic_trace.py:479 in call_module
ret_val = forward(*args, **kwargs)
File ~\anaconda3\Lib\site-packages\torch\fx\_symbolic_trace.py:788 in forward
return _orig_module_call(mod, *args, **kwargs)
File ~\anaconda3\Lib\site-packages\torch\nn\modules\module.py:1518 in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File ~\anaconda3\Lib\site-packages\torch\nn\modules\module.py:1527 in _call_impl
return forward_call(*args, **kwargs)
File ~\models\temporal_shift.py:21 in forward
return self.shift(x,self.n_segment,self.n_div)
File ~\models\temporal_shift.py:32 in shift
out[:,:-1,:fold]=x[:,1:,:fold]
TypeError: 'Proxy' object does not support item assignment
W0000 00:00:1787669093.121816 3408 face_landmarker_graph.cc:180] Sets FaceBlendshapesGraph acceleration to xnnpack by default.
W0000 00:00:1787669093.127475 33436 inference_feedback_manager.cc:121] Feedback manager requires a model with a single signature inference. Disabling support for feedback tensors.
W0000 00:00:1787669093.135861 33108 inference_feedback_manager.cc:121] Feedback manager requires a model with a single signature inference. Disabling support for feedback tensors.
the last bit is just mediapipe warning and can be ignored,basically model loads and dataloader is also fine but its the quant thats bugging.thing is when i ran the same quantization method get qconfig and then prepare_fx on simple resnet18 it worked then but now it fails.Any idea on how to fix it.
thanks dont want to rely on gpt based bugs as it halucinates a lot.