Spaces:
Runtime error
Runtime error
File size: 3,532 Bytes
f838ae8 2a75de9 7fe568e 2a75de9 f838ae8 f70f71e | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | from pytorch_lightning import LightningModule, Trainer, seed_everything
import torch
import pandas as pd
import numpy as np
import torch.nn as nn
from torch.nn import functional as F
from torchvision import transforms
from pytorch_lightning import LightningModule, Trainer, seed_everything
class LitResnet(LightningModule):
def __init__(self, lr=0.01, drop= 0.05, norm='BN',groupsize=1):
super().__init__()
self.save_hyperparameters()
self.num_classes =10
self.lr = lr
self.convblock1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=64, kernel_size=(3, 3), padding=1, bias=False),
nn.ReLU(),
self.user_norm(norm,64,groupsize),
nn.Dropout(drop),
nn.Conv2d(in_channels=64, out_channels=128, kernel_size=(3, 3), padding=1, bias=False),
nn.MaxPool2d(2,2),
nn.ReLU(),
self.user_norm(norm,128,groupsize),
nn.Dropout(drop))
self.res1 = nn.Sequential(
nn.Conv2d(in_channels=128, out_channels=128, kernel_size=(3, 3), padding=1, bias=False),
nn.ReLU(),
self.user_norm(norm,128,groupsize),
nn.Dropout(drop),
nn.Conv2d(in_channels = 128, out_channels=128, kernel_size=(3, 3), padding=1, bias=False),
nn.ReLU(),
self.user_norm(norm,128,groupsize),
nn.Dropout(drop)
)
# CONVOLUTION BLOCK 2
# Layer 2
self.convblock2 = nn.Sequential(
nn.Conv2d(in_channels=128, out_channels=256, kernel_size=(3, 3), padding=1, bias=False),
nn.ReLU(),
self.user_norm(norm,256,groupsize),
nn.Dropout(drop),
nn.MaxPool2d(2,2),
nn.ReLU(),
self.user_norm(norm,256,groupsize),
nn.Dropout(drop)
)
# CONVOLUTION BLOCK 3
self.convblock3 = nn.Sequential(
nn.Conv2d(in_channels=256, out_channels=512, kernel_size=(3, 3), padding=1, bias=False),
nn.MaxPool2d(2,2),
nn.ReLU(),
self.user_norm(norm,512,groupsize),
nn.Dropout(drop)
)
self.res2 = nn.Sequential(
nn.Conv2d(in_channels=512, out_channels=512, kernel_size=(3, 3), padding=1, bias=False),
nn.ReLU(),
self.user_norm(norm,512,groupsize),
nn.Dropout(drop),
nn.Conv2d(in_channels = 512, out_channels=512, kernel_size=(3, 3), padding=1, bias=False),
nn.ReLU(),
self.user_norm(norm,512,groupsize),
nn.Dropout(drop))
# CONVOLUTION BLOCK 4
self.convblock4 = nn.Sequential(
nn.MaxPool2d(kernel_size=4),
nn.Conv2d(in_channels=512, out_channels=10, kernel_size=(1, 1), padding=0, bias=False))
self.model = nn.Sequential(self.convblock1 ,self.convblock2, self.convblock3, self.convblock4)
def user_norm(self, norm, channels,groupsize=1):
if norm == 'BN':
return nn.BatchNorm2d(channels)
elif norm == 'LN':
return nn.GroupNorm(1,channels) #(equivalent with LayerNorm)
elif norm == 'GN':
return nn.GroupNorm(groupsize,channels) #groups=2
def forward(self, x):
x = self.convblock1(x)
x = x + self.res1 (x)
x = self.convblock2(x)
x = self.convblock3(x)
x = x + self.res2 (x)
x = self.convblock4(x)
x = x.view(-1, 10)
return F.log_softmax(x, dim=-1)
|