File size: 2,608 Bytes
c9f0165 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
from torch import nn
from torchvision import models
from torch.nn import *
import torch
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
class CustomResNet18(nn.Module):
def get_out_channels(self,module):
"""تابعی برای یافتن تعداد کانال‌های خروجی از لایه‌های کانولوشن و BatchNorm"""
if isinstance(module, nn.Conv2d):
return module.out_channels
elif isinstance(module, nn.BatchNorm2d):
return module.num_features
elif isinstance(module, nn.Linear):
return module.out_features
return None
def replace_relu_with_prelu_and_dropout(self,module, inplace=True):
for name, child in module.named_children():
# بازگشتی به لایه‌های زیرین
self.replace_relu_with_prelu_and_dropout(child, inplace)
if isinstance(child, nn.ReLU): # شناسایی لایه ReLU
# یافتن تعداد کانال‌های خروجی از ماژول قبلی
out_channels = None
for prev_name, prev_child in module.named_children():
if prev_name == name:
break
out_channels = self.get_out_channels(prev_child) or out_channels
if out_channels is None:
raise ValueError(f"Cannot determine `out_channels` for {child}. Please check the model structure.")
# ایجاد PReLU و Dropout2d
prelu = PReLU(device=device, num_parameters=out_channels) # استفاده از تعداد کانال‌های خروجی
dropout = nn.Dropout2d(p=0.2) # مقدار p تنظیم شده
# جایگزینی ReLU با Sequential شامل PReLU و Dropout
setattr(module, name, nn.Sequential(prelu, dropout).to(device))
def __init__(self):
super(CustomResNet18,self)
self.model = models.resnet18(weights = models.ResNet18_Weights.IMAGENET1K_V1).train(True).to(device)
self.replace_relu_with_prelu_and_dropout(self.model)
# print(model.fc.in_features)
number = self.model.fc.in_features
module = []
# استفاده از حلقه while برای تقسیم بر 2 تا رسیدن به عدد 8
module.append(LazyLinear(7))
self.model.fc = Sequential(*module).to(device)
def forward(self,x):
return self.model(x) |